mlir/test/Dialect/Affine/load-store.mlir

   1 // RUN: mlir-opt %s -split-input-file | FileCheck %s
   2
   3 // -----
   4
   5 // Test with just loop IVs.
   6 func.func @test0(%arg0 : index, %arg1 : index) {
   7   %0 = memref.alloc() : memref<100x100xf32>
   8   affine.for %i0 = 0 to 10 {
   9     affine.for %i1 = 0 to 10 {
  10       %1 = affine.load %0[%i0, %i1] : memref<100x100xf32>
  11 // CHECK: %{{.*}} = affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<100x100xf32>
  12     }
  13   }
  14   return
  15 }
  16
  17 // -----
  18
  19 // Test with loop IVs and constants.
  20 func.func @test1(%arg0 : index, %arg1 : index) {
  21   %0 = memref.alloc() : memref<100x100xf32>
  22   affine.for %i0 = 0 to 10 {
  23     affine.for %i1 = 0 to 10 {
  24       %1 = affine.load %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>
  25       affine.store %1, %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>
  26 // CHECK: %{{.*}} = affine.load %{{.*}}[%{{.*}} + 3, %{{.*}} + 7] : memref<100x100xf32>
  27 // CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}} + 3, %{{.*}} + 7] : memref<100x100xf32>
  28     }
  29   }
  30   return
  31 }
  32
  33 // -----
  34
  35 // Test with loop IVs and function args without 'symbol' keyword (should
  36 // be parsed as dim identifiers).
  37 func.func @test2(%arg0 : index, %arg1 : index) {
  38   %0 = memref.alloc() : memref<100x100xf32>
  39   affine.for %i0 = 0 to 10 {
  40     affine.for %i1 = 0 to 10 {
  41       %1 = affine.load %0[%i0 + %arg0, %i1 + %arg1] : memref<100x100xf32>
  42       affine.store %1, %0[%i0 + %arg0, %i1 + %arg1] : memref<100x100xf32>
  43 // CHECK: %{{.*}} = affine.load %{{.*}}[%{{.*}} + %{{.*}}, %{{.*}} + %{{.*}}] : memref<100x100xf32>
  44 // CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}} + %{{.*}}, %{{.*}} + %{{.*}}] : memref<100x100xf32>
  45     }
  46   }
  47   return
  48 }
  49
  50 // -----
  51
  52 // Test with loop IVs and function args with 'symbol' keyword (should
  53 // be parsed as symbol identifiers).
  54 func.func @test3(%arg0 : index, %arg1 : index) {
  55   %0 = memref.alloc() : memref<100x100xf32>
  56   affine.for %i0 = 0 to 10 {
  57     affine.for %i1 = 0 to 10 {
  58       %1 = affine.load %0[%i0 + symbol(%arg0), %i1 + symbol(%arg1)]
  59         : memref<100x100xf32>
  60       affine.store %1, %0[%i0 + symbol(%arg0), %i1 + symbol(%arg1)]
  61         : memref<100x100xf32>
  62 // CHECK: %{{.*}} = affine.load %{{.*}}[%{{.*}} + symbol(%{{.*}}), %{{.*}} + symbol(%{{.*}})] : memref<100x100xf32>
  63 // CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}} + symbol(%{{.*}}), %{{.*}} + symbol(%{{.*}})] : memref<100x100xf32>
  64     }
  65   }
  66   return
  67 }
  68
  69 // -----
  70
  71 // Test with loop IVs, symbols and constants in nested affine expressions.
  72 func.func @test4(%arg0 : index, %arg1 : index) {
  73   %0 = memref.alloc() : memref<100x100xf32>
  74   affine.for %i0 = 0 to 10 {
  75     affine.for %i1 = 0 to 10 {
  76       %1 = affine.load %0[(%i0 + symbol(%arg0)) floordiv 3 + 11,
  77                           (%i1 + symbol(%arg1)) mod 4 + 7] : memref<100x100xf32>
  78       affine.store %1, %0[(%i0 + symbol(%arg0)) floordiv 3 + 11,
  79                           (%i1 + symbol(%arg1)) mod 4 + 7] : memref<100x100xf32>
  80 // CHECK: %{{.*}} = affine.load %{{.*}}[(%{{.*}} + symbol(%{{.*}})) floordiv 3 + 11, (%{{.*}} + symbol(%{{.*}})) mod 4 + 7] : memref<100x100xf32>
  81 // CHECK: affine.store %{{.*}}, %{{.*}}[(%{{.*}} + symbol(%{{.*}})) floordiv 3 + 11, (%{{.*}} + symbol(%{{.*}})) mod 4 + 7] : memref<100x100xf32>
  82     }
  83   }
  84   return
  85 }
  86
  87 // -----
  88
  89 // Test with swizzled loop IVs.
  90 func.func @test5(%arg0 : index, %arg1 : index) {
  91   %0 = memref.alloc() : memref<10x10x10xf32>
  92   affine.for %i0 = 0 to 10 {
  93     affine.for %i1 = 0 to 10 {
  94       affine.for %i2 = 0 to 10 {
  95         %1 = affine.load %0[%i2, %i0, %i1] : memref<10x10x10xf32>
  96         affine.store %1, %0[%i2, %i0, %i1] : memref<10x10x10xf32>
  97 // CHECK: %{{.*}} = affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<10x10x10xf32>
  98 // CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] : memref<10x10x10xf32>
  99       }
 100     }
 101   }
 102   return
 103 }
 104
 105 // -----
 106
 107 // Test with swizzled loop IVs, duplicate args, and function args used as dims.
 108 // Dim identifiers are assigned in parse order:
 109 // d0 = %i2, d1 = %arg0, d2 = %i0, d3 = %i1, d4 = %arg1
 110 func.func @test6(%arg0 : index, %arg1 : index) {
 111   %0 = memref.alloc() : memref<10x10x10xf32>
 112   affine.for %i0 = 0 to 10 {
 113     affine.for %i1 = 0 to 10 {
 114       affine.for %i2 = 0 to 10 {
 115         %1 = affine.load %0[%i2 + %arg0, %i0 + %i1, %i1 + %arg0 + %arg1]
 116           : memref<10x10x10xf32>
 117         affine.store %1, %0[%i2 + %arg0, %i0 + %i1, %i1 + %arg0 + %arg1]
 118           : memref<10x10x10xf32>
 119 // CHECK: %{{.*}} = affine.load %{{.*}}[%{{.*}} + %{{.*}}, %{{.*}} + %{{.*}}, %{{.*}} + %{{.*}} + %{{.*}}] : memref<10x10x10xf32>
 120 // CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}} + %{{.*}}, %{{.*}} + %{{.*}}, %{{.*}} + %{{.*}} + %{{.*}}] : memref<10x10x10xf32>
 121       }
 122     }
 123   }
 124   return
 125 }
 126
 127 // -----
 128
 129 // Test with swizzled loop IVs, duplicate args, and function args used as syms.
 130 // Dim and symbol identifiers are assigned in parse order:
 131 // d0 = %i2, d1 = %i0, d2 = %i1
 132 // s0 = %arg0, s1 = %arg1
 133 func.func @test6(%arg0 : index, %arg1 : index) {
 134   %0 = memref.alloc() : memref<10x10x10xf32>
 135   affine.for %i0 = 0 to 10 {
 136     affine.for %i1 = 0 to 10 {
 137       affine.for %i2 = 0 to 10 {
 138         %1 = affine.load %0[%i2 + symbol(%arg0),
 139                             %i0 + %i1,
 140                             %i1 + symbol(%arg0) + symbol(%arg1)]
 141                               : memref<10x10x10xf32>
 142         affine.store %1, %0[%i2 + symbol(%arg0),
 143                              %i0 + %i1,
 144                              %i1 + symbol(%arg0) + symbol(%arg1)]
 145                               : memref<10x10x10xf32>
 146 // CHECK: %{{.*}} = affine.load %{{.*}}[%{{.*}} + symbol(%{{.*}}), %{{.*}} + %{{.*}}, %{{.*}} + symbol(%{{.*}}) + symbol(%{{.*}})] : memref<10x10x10xf32>
 147 // CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}} + symbol(%{{.*}}), %{{.*}} + %{{.*}}, %{{.*}} + symbol(%{{.*}}) + symbol(%{{.*}})] : memref<10x10x10xf32>
 148       }
 149     }
 150   }
 151   return
 152 }
 153
 154 // -----
 155
 156 // CHECK: [[MAP0:#map[0-9]*]] = affine_map<(d0) -> (d0 + 1)>
 157
 158 // Test with operands without special SSA name.
 159 func.func @test7() {
 160   %0 = memref.alloc() : memref<10xf32>
 161   affine.for %i0 = 0 to 10 {
 162     %1 = affine.apply affine_map<(d1) -> (d1 + 1)>(%i0)
 163     %2 = affine.load %0[%1] : memref<10xf32>
 164     affine.store %2, %0[%1] : memref<10xf32>
 165 // CHECK: affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
 166 // CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
 167   }
 168   return
 169 }
 170
 171 // -----
 172
 173 // Test with zero-dimensional operands.
 174 func.func @zero_dim(%arg0 : memref<i32>, %arg1 : memref<i32>) {
 175   %0 = affine.load %arg0[] : memref<i32>
 176   affine.store %0, %arg1[] : memref<i32>
 177   // CHECK: affine.load %{{.*}}[] : memref<i32>
 178   // CHECK: affine.store %{{.*}}, %{{.*}}[] : memref<i32>
 179   return
 180 }
 181
 182 // -----
 183
 184 // Test with loop IVs and constants.
 185 func.func @test_prefetch(%arg0 : index, %arg1 : index) {
 186   %0 = memref.alloc() : memref<100x100xf32>
 187   affine.for %i0 = 0 to 10 {
 188     affine.for %i1 = 0 to 10 {
 189       %1 = affine.load %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>
 190       affine.prefetch %0[%i0 + 3, %i1 + 11], write, locality<0>, data : memref<100x100xf32>
 191       // CHECK: affine.prefetch %{{.*}}[%{{.*}} + 3, %{{.*}} + 11], write, locality<0>, data : memref<100x100xf32>
 192       affine.prefetch %0[%i0, %i1 + 1], read, locality<3>, instr : memref<100x100xf32>
 193       // CHECK: affine.prefetch %{{.*}}[%{{.*}}, %{{.*}} + 1], read, locality<3>, instr : memref<100x100xf32>
 194     }
 195   }
 196   return
 197 }
 198
 199 // -----
 200
 201 // Test with just loop IVs.
 202 func.func @vector_load_vector_store_iv() {
 203   %0 = memref.alloc() : memref<100x100xf32>
 204   affine.for %i0 = 0 to 16 {
 205     affine.for %i1 = 0 to 16 step 8 {
 206       %1 = affine.vector_load %0[%i0, %i1] : memref<100x100xf32>, vector<8xf32>
 207       affine.vector_store %1, %0[%i0, %i1] : memref<100x100xf32>, vector<8xf32>
 208 // CHECK:      %[[buf:.*]] = memref.alloc
 209 // CHECK-NEXT: affine.for %[[i0:.*]] = 0
 210 // CHECK-NEXT:   affine.for %[[i1:.*]] = 0
 211 // CHECK-NEXT:     %[[val:.*]] = affine.vector_load %[[buf]][%[[i0]], %[[i1]]] : memref<100x100xf32>, vector<8xf32>
 212 // CHECK-NEXT:     affine.vector_store %[[val]], %[[buf]][%[[i0]], %[[i1]]] : memref<100x100xf32>, vector<8xf32>
 213     }
 214   }
 215   return
 216 }
 217
 218 // -----
 219
 220 // Test with loop IVs and constants.
 221 func.func @vector_load_vector_store_iv_constant() {
 222   %0 = memref.alloc() : memref<100x100xf32>
 223   affine.for %i0 = 0 to 10 {
 224     affine.for %i1 = 0 to 16 step 4 {
 225       %1 = affine.vector_load %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>, vector<4xf32>
 226       affine.vector_store %1, %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>, vector<4xf32>
 227 // CHECK:      %[[buf:.*]] = memref.alloc
 228 // CHECK-NEXT: affine.for %[[i0:.*]] = 0
 229 // CHECK-NEXT:   affine.for %[[i1:.*]] = 0
 230 // CHECK-NEXT:     %[[val:.*]] = affine.vector_load %{{.*}}[%{{.*}} + 3, %{{.*}} + 7] : memref<100x100xf32>, vector<4xf32>
 231 // CHECK-NEXT:     affine.vector_store %[[val]], %[[buf]][%[[i0]] + 3, %[[i1]] + 7] : memref<100x100xf32>, vector<4xf32>
 232     }
 233   }
 234   return
 235 }
 236
 237 // -----
 238
 239 func.func @vector_load_vector_store_2d() {
 240   %0 = memref.alloc() : memref<100x100xf32>
 241   affine.for %i0 = 0 to 16 step 2{
 242     affine.for %i1 = 0 to 16 step 8 {
 243       %1 = affine.vector_load %0[%i0, %i1] : memref<100x100xf32>, vector<2x8xf32>
 244       affine.vector_store %1, %0[%i0, %i1] : memref<100x100xf32>, vector<2x8xf32>
 245 // CHECK:      %[[buf:.*]] = memref.alloc
 246 // CHECK-NEXT: affine.for %[[i0:.*]] = 0
 247 // CHECK-NEXT:   affine.for %[[i1:.*]] = 0
 248 // CHECK-NEXT:     %[[val:.*]] = affine.vector_load %[[buf]][%[[i0]], %[[i1]]] : memref<100x100xf32>, vector<2x8xf32>
 249 // CHECK-NEXT:     affine.vector_store %[[val]], %[[buf]][%[[i0]], %[[i1]]] : memref<100x100xf32>, vector<2x8xf32>
 250     }
 251   }
 252   return
 253 }