mlir/test/Transforms/normalize-memrefs-ops.mlir

   1 // RUN: mlir-opt -normalize-memrefs %s | FileCheck %s
   2
   3 // For all these cases, we test if MemRefs Normalization works with the test
   4 // operations.
   5 // * test.op_norm: this operation has the MemRefsNormalizable attribute. The tests
   6 //   that include this operation are constructed so that the normalization should
   7 //   happen.
   8 // * test_op_nonnorm: this operation does not have the MemRefsNormalization
   9 //   attribute. The tests that include this operation are constructed so that the
  10 //    normalization should not happen.
  11
  12 #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 floordiv 32, d3 floordiv 64, d2 mod 32, d3 mod 64)>
  13
  14 // Test with op_norm and maps in arguments and in the operations in the function.
  15
  16 // CHECK-LABEL: test_norm
  17 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x64xf32>)
  18 func @test_norm(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
  19     %0 = alloc() : memref<1x16x14x14xf32, #map0>
  20     "test.op_norm"(%arg0, %0) : (memref<1x16x14x14xf32, #map0>, memref<1x16x14x14xf32, #map0>) -> ()
  21     dealloc %0 :  memref<1x16x14x14xf32, #map0>
  22
  23     // CHECK: %[[v0:.*]] = alloc() : memref<1x16x1x1x32x64xf32>
  24     // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x16x1x1x32x64xf32>, memref<1x16x1x1x32x64xf32>) -> ()
  25     // CHECK: dealloc %[[v0]] : memref<1x16x1x1x32x64xf32>
  26     return
  27 }
  28
  29 // Same test with op_nonnorm, with maps in the arguments and the operations in the function.
  30
  31 // CHECK-LABEL: test_nonnorm
  32 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x14x14xf32, #map>)
  33 func @test_nonnorm(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
  34     %0 = alloc() : memref<1x16x14x14xf32, #map0>
  35     "test.op_nonnorm"(%arg0, %0) : (memref<1x16x14x14xf32, #map0>, memref<1x16x14x14xf32, #map0>) -> ()
  36     dealloc %0 :  memref<1x16x14x14xf32, #map0>
  37
  38     // CHECK: %[[v0:.*]] = alloc() : memref<1x16x14x14xf32, #map>
  39     // CHECK: "test.op_nonnorm"(%[[ARG0]], %[[v0]]) : (memref<1x16x14x14xf32, #map>, memref<1x16x14x14xf32, #map>) -> ()
  40     // CHECK: dealloc %[[v0]] : memref<1x16x14x14xf32, #map>
  41     return
  42 }
  43
  44 // Test with op_norm, with maps in the operations in the function.
  45
  46 // CHECK-LABEL: test_norm_mix
  47 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x64xf32>
  48 func @test_norm_mix(%arg0 : memref<1x16x1x1x32x64xf32>) -> () {
  49     %0 = alloc() : memref<1x16x14x14xf32, #map0>
  50     "test.op_norm"(%arg0, %0) : (memref<1x16x1x1x32x64xf32>, memref<1x16x14x14xf32, #map0>) -> ()
  51     dealloc %0 :  memref<1x16x14x14xf32, #map0>
  52
  53     // CHECK: %[[v0:.*]] = alloc() : memref<1x16x1x1x32x64xf32>
  54     // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x16x1x1x32x64xf32>, memref<1x16x1x1x32x64xf32>) -> ()
  55     // CHECK: dealloc %[[v0]] : memref<1x16x1x1x32x64xf32>
  56     return
  57 }
  58
  59 // Test with maps in load and store ops.
  60
  61 #map_tile = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2 floordiv 32, d3 floordiv 32, d2 mod 32, d3 mod 32)>
  62
  63 // CHECK-LABEL: test_load_store
  64 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x14x14xf32>
  65 func @test_load_store(%arg0 : memref<1x16x14x14xf32>) -> () {
  66     %0 = alloc() : memref<1x16x14x14xf32, #map_tile>
  67     // CHECK: %[[v0:.*]] = alloc() : memref<1x16x1x1x32x32xf32>
  68     %1 = alloc() : memref<1x16x14x14xf32>
  69     // CHECK: %[[v1:.*]] = alloc() : memref<1x16x14x14xf32>
  70     "test.op_norm"(%0, %1) : (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>) -> ()
  71     // CHECK: "test.op_norm"(%[[v0]], %[[v1]]) : (memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>) -> ()
  72     %cst = constant 3.0 : f32
  73     affine.for %i = 0 to 1 {
  74       affine.for %j = 0 to 16 {
  75         affine.for %k = 0 to 14 {
  76           affine.for %l = 0 to 14 {
  77             %2 = load %1[%i, %j, %k, %l] : memref<1x16x14x14xf32>
  78             // CHECK: memref<1x16x14x14xf32>
  79             %3 = addf %2, %cst : f32
  80             store %3, %arg0[%i, %j, %k, %l] : memref<1x16x14x14xf32>
  81             // CHECK: memref<1x16x14x14xf32>
  82           }
  83         }
  84       }
  85     }
  86     dealloc %0 :  memref<1x16x14x14xf32, #map_tile>
  87     // CHECK: dealloc %[[v0]] : memref<1x16x1x1x32x32xf32>
  88     dealloc %1 :  memref<1x16x14x14xf32>
  89     // CHECK: dealloc %[[v1]] : memref<1x16x14x14xf32>
  90     return
  91 }
  92
  93 // Test with op_norm_ret, with maps in the results of normalizable operation.
  94
  95 // CHECK-LABEL: test_norm_ret
  96 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x32xf32>) -> (memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>) {
  97 func @test_norm_ret(%arg0: memref<1x16x14x14xf32, #map_tile>) -> (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>) {
  98     %0 = alloc() : memref<1x16x14x14xf32, #map_tile>
  99     // CHECK-NEXT: %[[v0:.*]] = alloc() : memref<1x16x1x1x32x32xf32>
 100     %1, %2 = "test.op_norm_ret"(%arg0) : (memref<1x16x14x14xf32, #map_tile>) -> (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>)
 101     // CHECK-NEXT: %[[v1:.*]], %[[v2:.*]] = "test.op_norm_ret"
 102     // CHECK-SAME: (memref<1x16x1x1x32x32xf32>) -> (memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>)
 103     "test.op_norm"(%1, %0) : (memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32, #map_tile>) -> ()
 104     // CHECK-NEXT: "test.op_norm"
 105     // CHECK-SAME: : (memref<1x16x1x1x32x32xf32>, memref<1x16x1x1x32x32xf32>) -> ()
 106     dealloc %0 : memref<1x16x14x14xf32, #map_tile>
 107     // CHECK-NEXT: dealloc %[[v0]] : memref<1x16x1x1x32x32xf32>
 108     return %1, %2 : memref<1x16x14x14xf32, #map_tile>, memref<1x16x14x14xf32>
 109     // CHECK-NEXT: return %[[v1]], %[[v2]] : memref<1x16x1x1x32x32xf32>, memref<1x16x14x14xf32>
 110 }
 111
 112 // Test with an arbitrary op that references the function symbol.
 113
 114 "test.op_funcref"() {func = @test_norm_mix} : () -> ()