254 lines · plain
1// RUN: mlir-opt %s -split-input-file | FileCheck %s2 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 return15}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 return31}32 33// -----34 35// Test with loop IVs and function args without 'symbol' keyword (should36// 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 return48}49 50// -----51 52// Test with loop IVs and function args with 'symbol' keyword (should53// 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 return67}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 return85}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 return103}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 = %arg1110func.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 return125}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 = %i1132// s0 = %arg0, s1 = %arg1133func.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 return152}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 return169}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 return180}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 return197}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.alloc209// CHECK-NEXT: affine.for %[[i0:.*]] = 0210// CHECK-NEXT: affine.for %[[i1:.*]] = 0211// 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 return216}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.alloc228// CHECK-NEXT: affine.for %[[i0:.*]] = 0229// CHECK-NEXT: affine.for %[[i1:.*]] = 0230// 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 return235}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.alloc246// CHECK-NEXT: affine.for %[[i0:.*]] = 0247// CHECK-NEXT: affine.for %[[i1:.*]] = 0248// 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 return253}254