mlir/test/Dialect/Affine/memref-stride-calculation.mlir

   1 // RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(test-memref-stride-calculation))" -o /dev/null | FileCheck %s
   2
   3 func.func @f(%0: index) {
   4 // CHECK-LABEL: Testing: f
   5   %1 = memref.alloc() : memref<3x4x5xf32>
   6 // CHECK: MemRefType offset: 0 strides: 20, 5, 1
   7   %2 = memref.alloc(%0) : memref<3x4x?xf32>
   8 // CHECK: MemRefType offset: 0 strides: ?, ?, 1
   9   %3 = memref.alloc(%0) : memref<3x?x5xf32>
  10 // CHECK: MemRefType offset: 0 strides: ?, 5, 1
  11   %4 = memref.alloc(%0) : memref<?x4x5xf32>
  12 // CHECK: MemRefType offset: 0 strides: 20, 5, 1
  13   %5 = memref.alloc(%0, %0) : memref<?x4x?xf32>
  14 // CHECK: MemRefType offset: 0 strides: ?, ?, 1
  15   %6 = memref.alloc(%0, %0, %0) : memref<?x?x?xf32>
  16 // CHECK: MemRefType offset: 0 strides: ?, ?, 1
  17
  18   %11 = memref.alloc() : memref<3x4x5xf32, affine_map<(i, j, k)->(i, j, k)>>
  19 // CHECK: MemRefType offset: 0 strides: 20, 5, 1
  20   %b11 = memref.alloc() : memref<3x4x5xf32, strided<[20, 5, 1], offset: 0>>
  21 // CHECK: MemRefType offset: 0 strides: 20, 5, 1
  22   %12 = memref.alloc(%0) : memref<3x4x?xf32, affine_map<(i, j, k)->(i, j, k)>>
  23 // CHECK: MemRefType offset: 0 strides: ?, ?, 1
  24   %13 = memref.alloc(%0) : memref<3x?x5xf32, affine_map<(i, j, k)->(i, j, k)>>
  25 // CHECK: MemRefType offset: 0 strides: ?, 5, 1
  26   %14 = memref.alloc(%0) : memref<?x4x5xf32, affine_map<(i, j, k)->(i, j, k)>>
  27 // CHECK: MemRefType offset: 0 strides: 20, 5, 1
  28   %15 = memref.alloc(%0, %0) : memref<?x4x?xf32, affine_map<(i, j, k)->(i, j, k)>>
  29 // CHECK: MemRefType offset: 0 strides: ?, ?, 1
  30   %16 = memref.alloc(%0, %0, %0) : memref<?x?x?xf32, affine_map<(i, j, k)->(i, j, k)>>
  31 // CHECK: MemRefType offset: 0 strides: ?, ?, 1
  32
  33   %21 = memref.alloc()[%0] : memref<3x4x5xf32, affine_map<(i, j, k)[M]->(32 * i + 16 * j + M * k + 1)>>
  34 // CHECK: MemRefType offset: 1 strides: 32, 16, ?
  35   %22 = memref.alloc()[%0] : memref<3x4x5xf32, affine_map<(i, j, k)[M]->(32 * i + M * j + 16 * k + 3)>>
  36 // CHECK: MemRefType offset: 3 strides: 32, ?, 16
  37   %b22 = memref.alloc(%0)[%0, %0] : memref<3x4x?xf32, strided<[?, ?, 1], offset: 0>>
  38 // CHECK: MemRefType offset: 0 strides: ?, ?, 1
  39   %23 = memref.alloc(%0)[%0] : memref<3x?x5xf32, affine_map<(i, j, k)[M]->(M * i + 32 * j + 16 * k + 7)>>
  40 // CHECK: MemRefType offset: 7 strides: ?, 32, 16
  41   %b23 = memref.alloc(%0)[%0] : memref<3x?x5xf32, strided<[?, 5, 1], offset: 0>>
  42 // CHECK: MemRefType offset: 0 strides: ?, 5, 1
  43   %24 = memref.alloc(%0)[%0] : memref<3x?x5xf32, affine_map<(i, j, k)[M]->(M * i + 32 * j + 16 * k + M)>>
  44 // CHECK: MemRefType offset: ? strides: ?, 32, 16
  45   %b24 = memref.alloc(%0)[%0, %0] : memref<3x?x5xf32, strided<[?, 32, 16], offset: ?>>
  46 // CHECK: MemRefType offset: ? strides: ?, 32, 16
  47   %25 = memref.alloc(%0, %0)[%0, %0] : memref<?x?x16xf32, affine_map<(i, j, k)[M, N]->(M * i + N * j + k + 1)>>
  48 // CHECK: MemRefType offset: 1 strides: ?, ?, 1
  49   %b25 = memref.alloc(%0, %0)[%0, %0] : memref<?x?x16xf32, strided<[?, ?, 1], offset: 1>>
  50 // CHECK: MemRefType offset: 1 strides: ?, ?, 1
  51   %26 = memref.alloc(%0)[] : memref<?xf32, affine_map<(i)[M]->(i)>>
  52 // CHECK: MemRefType offset: 0 strides: 1
  53   %27 = memref.alloc()[%0] : memref<5xf32, affine_map<(i)[M]->(M)>>
  54 // CHECK: MemRefType memref<5xf32, affine_map<(d0)[s0] -> (s0)>> cannot be converted to strided form
  55   %28 = memref.alloc()[%0] : memref<5xf32, affine_map<(i)[M]->(123)>>
  56 // CHECK: MemRefType memref<5xf32, affine_map<(d0)[s0] -> (123)>> cannot be converted to strided form
  57   %29 = memref.alloc()[%0] : memref<f32, affine_map<()[M]->(M)>>
  58 // CHECK: MemRefType offset: ? strides:
  59   %30 = memref.alloc()[%0] : memref<f32, affine_map<()[M]->(123)>>
  60 // CHECK: MemRefType offset: 123 strides:
  61
  62   %101 = memref.alloc() : memref<3x4x5xf32, affine_map<(i, j, k)->(i floordiv 4 + j + k)>>
  63 // CHECK: MemRefType memref<3x4x5xf32, affine_map<(d0, d1, d2) -> (d0 floordiv 4 + d1 + d2)>> cannot be converted to strided form
  64   %102 = memref.alloc() : memref<3x4x5xf32, affine_map<(i, j, k)->(i ceildiv 4 + j + k)>>
  65 // CHECK: MemRefType memref<3x4x5xf32, affine_map<(d0, d1, d2) -> (d0 ceildiv 4 + d1 + d2)>> cannot be converted to strided form
  66   %103 = memref.alloc() : memref<3x4x5xf32, affine_map<(i, j, k)->(i mod 4 + j + k)>>
  67 // CHECK: MemRefType memref<3x4x5xf32, affine_map<(d0, d1, d2) -> (d0 mod 4 + d1 + d2)>> cannot be converted to strided form
  68
  69   %200 = memref.alloc()[%0, %0, %0] : memref<3x4x5xf32, affine_map<(i, j, k)[M, N, K]->(M * i + N * i + N * j + K * k - (M + N - 20)* i)>>
  70   // CHECK: MemRefType offset: 0 strides: 20, ?, ?
  71   %201 = memref.alloc()[%0, %0, %0] : memref<3x4x5xf32, affine_map<(i, j, k)[M, N, K]->(M * i + N * i + N * K * j + K * K * k - (M + N - 20) * (i + 1))>>
  72   // CHECK: MemRefType offset: ? strides: 20, ?, ?
  73   %202 = memref.alloc()[%0, %0, %0] : memref<3x4x5xf32, affine_map<(i, j, k)[M, N, K]->(M * (i + 1) + j + k - M)>>
  74   // CHECK: MemRefType offset: 0 strides: ?, 1, 1
  75   %203 = memref.alloc()[%0, %0, %0] : memref<3x4x5xf32, affine_map<(i, j, k)[M, N, K]->(M + M * (i + N * (j + K * k)))>>
  76   // CHECK: MemRefType offset: ? strides: ?, ?, ?
  77
  78   return
  79 }