429 lines · plain
1// RUN: mlir-opt -allow-unregistered-dialect -split-input-file %s | FileCheck %s2// RUN: mlir-opt -allow-unregistered-dialect -split-input-file %s -mlir-print-op-generic | FileCheck -check-prefix=GENERIC %s3 4// Check that the attributes for the affine operations are round-tripped.5// Check that `affine.yield` is visible in the generic form.6// CHECK-LABEL: @empty7func.func @empty() {8 // CHECK: affine.for9 // CHECK-NEXT: } {some_attr = true}10 //11 // GENERIC: "affine.for"()12 // GENERIC-NEXT: ^bb0(%{{.*}}: index):13 // GENERIC-NEXT: "affine.yield"() : () -> ()14 // GENERIC-NEXT: })15 affine.for %i = 0 to 10 {16 } {some_attr = true}17 18 // CHECK: affine.if19 // CHECK-NEXT: } {some_attr = true}20 //21 // GENERIC: "affine.if"()22 // GENERIC-NEXT: "affine.yield"() : () -> ()23 // GENERIC-NEXT: }, {24 // GENERIC-NEXT: })25 affine.if affine_set<() : ()> () {26 } {some_attr = true}27 28 // CHECK: } else {29 // CHECK: } {some_attr = true}30 //31 // GENERIC: "affine.if"()32 // GENERIC-NEXT: "affine.yield"() : () -> ()33 // GENERIC-NEXT: }, {34 // GENERIC-NEXT: "foo"() : () -> ()35 // GENERIC-NEXT: "affine.yield"() : () -> ()36 // GENERIC-NEXT: })37 affine.if affine_set<() : ()> () {38 } else {39 "foo"() : () -> ()40 } {some_attr = true}41 42 return43}44 45// -----46 47// GENERIC: #[[$map:.*]] = affine_map<() -> (0)>48// GENERIC: #[[$map1:.*]] = affine_map<() -> (10)>49 50// Check that an explicit affine.yield is not printed in custom format.51// Check that no extra terminator is introduced.52// CHECK-LABEL: @affine.yield53// CHECK-GENERIC-LABEL: @affine.yield54func.func @affine.yield() {55 // CHECK: affine.for56 // CHECK-NEXT: }57 //58 // GENERIC: "affine.for"() <{lowerBoundMap = #[[$map]], operandSegmentSizes = array<i32: 0, 0, 0>, step = 1 : index, upperBoundMap = #[[$map1]]}> ({59 // GENERIC-NEXT: ^bb0(%{{.*}}: index):60 // GENERIC-NEXT: "affine.yield"() : () -> ()61 // GENERIC-NEXT: }) : () -> ()62 affine.for %i = 0 to 10 {63 "affine.yield"() : () -> ()64 }65 return66}67 68// -----69 70// CHECK-DAG: #[[$MAP0:map[0-9]*]] = affine_map<(d0)[s0] -> (1000, d0 + 512, s0)>71// CHECK-DAG: #[[$MAP1:map[0-9]*]] = affine_map<(d0, d1)[s0] -> (d0 - d1, s0 + 512)>72// CHECK-DAG: #[[$MAP2:map[0-9]*]] = affine_map<()[s0, s1] -> (s0 - s1, 11)>73// CHECK-DAG: #[[$MAP3:map[0-9]*]] = affine_map<() -> (77, 78, 79)>74 75// CHECK-LABEL: @affine_min76func.func @affine_min(%arg0 : index, %arg1 : index, %arg2 : index) {77 // CHECK: affine.min #[[$MAP0]](%arg0)[%arg1]78 %0 = affine.min affine_map<(d0)[s0] -> (1000, d0 + 512, s0)> (%arg0)[%arg1]79 // CHECK: affine.min #[[$MAP1]](%arg0, %arg1)[%arg2]80 %1 = affine.min affine_map<(d0, d1)[s0] -> (d0 - d1, s0 + 512)> (%arg0, %arg1)[%arg2]81 // CHECK: affine.min #[[$MAP2]]()[%arg1, %arg2]82 %2 = affine.min affine_map<()[s0, s1] -> (s0 - s1, 11)> ()[%arg1, %arg2]83 // CHECK: affine.min #[[$MAP3]]()84 %3 = affine.min affine_map<()[] -> (77, 78, 79)> ()[]85 return86}87 88// CHECK-LABEL: @affine_max89func.func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {90 // CHECK: affine.max #[[$MAP0]](%arg0)[%arg1]91 %0 = affine.max affine_map<(d0)[s0] -> (1000, d0 + 512, s0)> (%arg0)[%arg1]92 // CHECK: affine.max #[[$MAP1]](%arg0, %arg1)[%arg2]93 %1 = affine.max affine_map<(d0, d1)[s0] -> (d0 - d1, s0 + 512)> (%arg0, %arg1)[%arg2]94 // CHECK: affine.max #[[$MAP2]]()[%arg1, %arg2]95 %2 = affine.max affine_map<()[s0, s1] -> (s0 - s1, 11)> ()[%arg1, %arg2]96 // CHECK: affine.max #[[$MAP3]]()97 %3 = affine.max affine_map<()[] -> (77, 78, 79)> ()[]98 return99}100 101// -----102 103func.func @valid_symbols(%arg0: index, %arg1: index, %arg2: index) {104 %c1 = arith.constant 1 : index105 %c0 = arith.constant 0 : index106 %0 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>107 affine.for %arg3 = 0 to %arg2 step 768 {108 %13 = memref.dim %0, %c1 : memref<?x?xf32>109 affine.for %arg4 = 0 to %13 step 264 {110 %18 = memref.dim %0, %c0 : memref<?x?xf32>111 %20 = memref.subview %0[%c0, %c0][%18,%arg4][%c1,%c1] : memref<?x?xf32>112 to memref<?x?xf32, strided<[?, ?], offset: ?>>113 %24 = memref.dim %20, %c0 : memref<?x?xf32, strided<[?, ?], offset: ?>>114 affine.for %arg5 = 0 to %24 step 768 {115 "foo"() : () -> ()116 }117 }118 }119 return120}121 122// -----123 124// Test symbol constraints for ops with AffineScope trait.125 126// CHECK-LABEL: func @valid_symbol_affine_scope127func.func @valid_symbol_affine_scope(%n : index, %A : memref<?xf32>) {128 test.affine_scope {129 %c1 = arith.constant 1 : index130 %l = arith.subi %n, %c1 : index131 // %l, %n are valid symbols since test.affine_scope defines a new affine132 // scope.133 affine.for %i = %l to %n {134 %m = arith.subi %l, %i : index135 test.affine_scope {136 // %m and %n are valid symbols.137 affine.for %j = %m to %n {138 %v = affine.load %A[%n - 1] : memref<?xf32>139 affine.store %v, %A[%n - 1] : memref<?xf32>140 }141 "terminate"() : () -> ()142 }143 }144 "terminate"() : () -> ()145 }146 return147}148 149// -----150 151// Test dimension constraints for linearize_index and delinearize_index152 153// CHECK-LABEL: func @valid_dim_linearize_delinearize154func.func @valid_dim_linearize_delinearize(%m : index, %n : index, %A : memref<?xf32>, %B: memref<?x32x?xf32>) {155 affine.for %0 = 0 to %m {156 affine.for %1 = 0 to %n {157 %load_idx = affine.linearize_index disjoint [%0, %1] by (%m, %n) : index158 %store_idx0, %store_idx1 = affine.delinearize_index %n into (32) : index, index159 %v = affine.load %A[%load_idx] : memref<?xf32>160 affine.store %v, %B[%0, %store_idx1, %store_idx0] : memref<?x32x?xf32>161 }162 }163 return164}165 166// -----167 168// Test the fact that module op always provides an affine scope.169 170%idx = "test.foo"() : () -> (index)171"test.func"() ({172^bb0(%A : memref<?xf32>):173 affine.load %A[%idx] : memref<?xf32>174 "terminate"() : () -> ()175}) : () -> ()176 177// -----178 179// CHECK-LABEL: func @parallel180// CHECK-SAME: (%[[A:.*]]: memref<100x100xf32>, %[[N:.*]]: index)181func.func @parallel(%A : memref<100x100xf32>, %N : index) {182 // CHECK: affine.parallel (%[[I0:.*]], %[[J0:.*]]) = (0, 0) to (symbol(%[[N]]), 100) step (10, 10)183 affine.parallel (%i0, %j0) = (0, 0) to (symbol(%N), 100) step (10, 10) {184 // CHECK: affine.parallel (%{{.*}}, %{{.*}}) = (%[[I0]], %[[J0]]) to (%[[I0]] + 10, %[[J0]] + 10) reduce ("minimumf", "maximumf") -> (f32, f32)185 %0:2 = affine.parallel (%i1, %j1) = (%i0, %j0) to (%i0 + 10, %j0 + 10) reduce ("minimumf", "maximumf") -> (f32, f32) {186 %2 = affine.load %A[%i0 + %i0, %j0 + %j1] : memref<100x100xf32>187 affine.yield %2, %2 : f32, f32188 }189 }190 return191}192 193// -----194 195// CHECK-LABEL: @parallel_min_max196// CHECK: %[[A:.*]]: index, %[[B:.*]]: index, %[[C:.*]]: index, %[[D:.*]]: index197func.func @parallel_min_max(%a: index, %b: index, %c: index, %d: index) {198 // CHECK: affine.parallel (%{{.*}}, %{{.*}}, %{{.*}}) =199 // CHECK: (max(%[[A]], %[[B]])200 // CHECK: to (%[[C]], min(%[[C]], %[[D]]), %[[B]])201 affine.parallel (%i, %j, %k) = (max(%a, %b), %b, max(%a, %c))202 to (%c, min(%c, %d), %b) {203 affine.yield204 }205 return206}207 208// -----209 210// CHECK-LABEL: @parallel_no_ivs211func.func @parallel_no_ivs() {212 // CHECK: affine.parallel () = () to ()213 affine.parallel () = () to () {214 affine.yield215 }216 return217}218 219// -----220 221// CHECK-LABEL: func @affine_if222func.func @affine_if() -> f32 {223 // CHECK: %[[ZERO:.*]] = arith.constant {{.*}} : f32224 %zero = arith.constant 0.0 : f32225 // CHECK: %[[OUT:.*]] = affine.if {{.*}}() -> f32 {226 %0 = affine.if affine_set<() : ()> () -> f32 {227 // CHECK: affine.yield %[[ZERO]] : f32228 affine.yield %zero : f32229 } else {230 // CHECK: affine.yield %[[ZERO]] : f32231 affine.yield %zero : f32232 }233 // CHECK: return %[[OUT]] : f32234 return %0 : f32235}236 237// -----238 239// Test affine.for with yield values.240 241#set = affine_set<(d0): (d0 - 10 >= 0)>242 243// CHECK-LABEL: func @yield_loop244func.func @yield_loop(%buffer: memref<1024xf32>) -> f32 {245 %sum_init_0 = arith.constant 0.0 : f32246 %res = affine.for %i = 0 to 10 step 2 iter_args(%sum_iter = %sum_init_0) -> f32 {247 %t = affine.load %buffer[%i] : memref<1024xf32>248 %sum_next = affine.if #set(%i) -> (f32) {249 %new_sum = arith.addf %sum_iter, %t : f32250 affine.yield %new_sum : f32251 } else {252 affine.yield %sum_iter : f32253 }254 affine.yield %sum_next : f32255 }256 return %res : f32257}258// CHECK: %[[const_0:.*]] = arith.constant 0.000000e+00 : f32259// CHECK-NEXT: %[[output:.*]] = affine.for %{{.*}} = 0 to 10 step 2 iter_args(%{{.*}} = %[[const_0]]) -> (f32) {260// CHECK: affine.if #set(%{{.*}}) -> f32 {261// CHECK: affine.yield %{{.*}} : f32262// CHECK-NEXT: } else {263// CHECK-NEXT: affine.yield %{{.*}} : f32264// CHECK-NEXT: }265// CHECK-NEXT: affine.yield %{{.*}} : f32266// CHECK-NEXT: }267// CHECK-NEXT: return %[[output]] : f32268 269// CHECK-LABEL: func @affine_for_multiple_yield270func.func @affine_for_multiple_yield(%buffer: memref<1024xf32>) -> (f32, f32) {271 %init_0 = arith.constant 0.0 : f32272 %res1, %res2 = affine.for %i = 0 to 10 step 2 iter_args(%iter_arg1 = %init_0, %iter_arg2 = %init_0) -> (f32, f32) {273 %t = affine.load %buffer[%i] : memref<1024xf32>274 %ret1 = arith.addf %t, %iter_arg1 : f32275 %ret2 = arith.addf %t, %iter_arg2 : f32276 affine.yield %ret1, %ret2 : f32, f32277 }278 return %res1, %res2 : f32, f32279}280// CHECK: %[[const_0:.*]] = arith.constant 0.000000e+00 : f32281// CHECK-NEXT: %[[output:[0-9]+]]:2 = affine.for %{{.*}} = 0 to 10 step 2 iter_args(%[[iter_arg1:.*]] = %[[const_0]], %[[iter_arg2:.*]] = %[[const_0]]) -> (f32, f32) {282// CHECK: %[[res1:.*]] = arith.addf %{{.*}}, %[[iter_arg1]] : f32283// CHECK-NEXT: %[[res2:.*]] = arith.addf %{{.*}}, %[[iter_arg2]] : f32284// CHECK-NEXT: affine.yield %[[res1]], %[[res2]] : f32, f32285// CHECK-NEXT: }286 287// -----288 289// CHECK-LABEL: func @delinearize290func.func @delinearize(%linear_idx: index, %basis0: index, %basis1 :index) -> (index, index) {291 // CHECK: affine.delinearize_index %{{.+}} into (%{{.+}}, %{{.+}}) : index, index292 %1:2 = affine.delinearize_index %linear_idx into (%basis0, %basis1) : index, index293 return %1#0, %1#1 : index, index294}295 296// CHECK-LABEL: @delinearize_mixed297func.func @delinearize_mixed(%linear_idx: index, %basis1: index) -> (index, index, index) {298 // CHECK: affine.delinearize_index %{{.+}} into (2, %{{.+}}, 3) : index, index, index299 %1:3 = affine.delinearize_index %linear_idx into (2, %basis1, 3) : index, index, index300 return %1#0, %1#1, %1#2 : index, index, index301}302 303// -----304 305// CHECK-LABEL: func @linearize306func.func @linearize(%index0: index, %index1: index, %basis0: index, %basis1 :index) -> index {307 // CHECK: affine.linearize_index [%{{.+}}, %{{.+}}] by (%{{.+}}, %{{.+}}) : index308 %1 = affine.linearize_index [%index0, %index1] by (%basis0, %basis1) : index309 return %1 : index310}311 312// CHECK-LABEL: @linearize_mixed313func.func @linearize_mixed(%index0: index, %index1: index, %index2: index, %basis1: index) -> index {314 // CHECK: affine.linearize_index disjoint [%{{.+}}, %{{.+}}, %{{.+}}] by (2, %{{.+}}, 3) : index315 %1 = affine.linearize_index disjoint [%index0, %index1, %index2] by (2, %basis1, 3) : index316 return %1 : index317}318 319// -----320 321// CHECK-LABEL: @gpu_launch_affine322 323// Test `thread_id` in AffineScope, the `thread_id` is in AffineScope's toplevel,324// it is a valid symbol.325 326module {327 func.func @gpu_launch_affine() {328 %c1 = arith.constant 1 : index329 gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1)330 threads(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) {331 %thread_id_x = gpu.thread_id x332 %c128 = arith.constant 128 : index333 affine.for %arg12 = %thread_id_x to %c128 step 8 {334 }335 gpu.terminator336 }337 return338 }339}340 341// CHECK: %[[THREAD_ID:.*]] = gpu.thread_id x342// CHECK: %[[VAL:.*]] = arith.constant 128 : index343// CHECK: affine.for %{{.*}} = %[[THREAD_ID]] to %[[VAL]] step 8 {344 345// -----346 347#map = affine_map<()[s0] -> (s0 mod 32)>348 349// CHECK: #[[$ATTR_0:.+]] = affine_map<()[s0] -> (s0 mod 32)>350 351// CHECK-LABEL: gpu.func @affine_thread_id352 353module {354 gpu.module @gpu {355 gpu.func @affine_thread_id(%arg0: memref<?x?xf32>) kernel {356 %c3 = arith.constant 3 : index357 %dim = memref.dim %arg0, %c3 : memref<?x?xf32>358 %c0 = arith.constant 0 : index359 affine.for %arg3 = %c0 to %dim step 32 {360 %thread_id_x = gpu.thread_id x361 %0 = affine.apply #map()[%thread_id_x]362 %c128 = arith.constant 128 : index363 affine.for %arg4 = %0 to %c128 step 8 {364 %c32 = arith.constant 32 : index365 }366 }367 gpu.return368 }369 }370}371 372// CHECK-SAME: (%[[VAL_0:.*]]: memref<?x?xf32>) kernel {373// CHECK: %[[VAL_1:.*]] = arith.constant 3 : index374// CHECK: %[[VAL_2:.*]] = memref.dim %[[VAL_0]], %[[VAL_1]] : memref<?x?xf32>375// CHECK: %[[VAL_3:.*]] = arith.constant 0 : index376// CHECK: affine.for %[[VAL_4:.*]] = %[[VAL_3]] to %[[VAL_2]] step 32 {377// CHECK: %[[VAL_5:.*]] = gpu.thread_id x378// CHECK: %[[VAL_6:.*]] = affine.apply #[[$ATTR_0]](){{\[}}%[[VAL_5]]]379// CHECK: %[[VAL_7:.*]] = arith.constant 128 : index380// CHECK: affine.for %{{.*}} = %[[VAL_6]] to %[[VAL_7]] step 8 {381 382// -----383 384#map = affine_map<(d0)[s0] -> (d0 + s0)>385 386// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>387 388// CHECK-LABEL: func @arith_add_vaild_symbol_upper_bound389 390func.func @arith_add_vaild_symbol_upper_bound(%arg : index) {391 affine.for %n0 = 0 to 7 {392 %dim = arith.addi %arg, %arg : index393 affine.for %n1 = 0 to #map(%dim)[%arg] {394 }395 }396 return397}398 399// CHECK-SAME: %[[VAL_0:.*]]: index) {400// CHECK: affine.for %[[VAL_1:.*]] = 0 to 7 {401// CHECK: %[[VAL_2:.*]] = arith.addi %[[VAL_0]], %[[VAL_0]] : index402// CHECK: affine.for %[[VAL_3:.*]] = 0 to #[[$ATTR_0]](%[[VAL_2]]){{\[}}%[[VAL_0]]] {403// CHECK: }404// CHECK: }405 406// -----407 408#map = affine_map<(d0)[s0] -> (d0 + s0)>409 410// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>411 412// CHECK-LABEL: func @arith_add_vaild_symbol_lower_bound413 414func.func @arith_add_vaild_symbol_lower_bound(%arg : index) {415 affine.for %n0 = 0 to 7 {416 %dim = arith.addi %arg, %arg : index417 affine.for %n1 = #map(%dim)[%arg] to 7 {418 }419 }420 return421}422 423// CHECK-SAME: %[[VAL_0:.*]]: index) {424// CHECK: affine.for %[[VAL_1:.*]] = 0 to 7 {425// CHECK: %[[VAL_2:.*]] = arith.addi %[[VAL_0]], %[[VAL_0]] : index426// CHECK: affine.for %[[VAL_3:.*]] = #[[$ATTR_0]](%[[VAL_2]]){{\[}}%[[VAL_0]]] to 7 {427// CHECK: }428// CHECK: }429