1// RUN: mlir-opt %s -split-input-file | FileCheck %s 2 3// ----- 4 5// Test with just loop IVs. 6func.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. 20func.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). 37func.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). 54func.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. 72func.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. 90func.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 110func.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 133func.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. 159func.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. 174func.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. 185func.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. 202func.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. 221func.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 239func.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} 254