553 lines · plain
1// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll-jam{unroll-jam-factor=2}))" | FileCheck %s2// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll-jam{unroll-jam-factor=4}))" | FileCheck --check-prefix=UJAM-FOUR %s3// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(gpu.module(gpu.func(affine-loop-unroll-jam{unroll-jam-factor=2})))" | FileCheck --check-prefix=GPU-HJAM %s4 5// CHECK-DAG: [[$MAP_PLUS_1:#map[0-9]*]] = affine_map<(d0) -> (d0 + 1)>6// CHECK-DAG: [[$MAP_DIV_OFFSET:#map[0-9]*]] = affine_map<()[s0] -> (((s0 - 1) floordiv 2) * 2 + 1)>7// CHECK-DAG: [[$MAP_SYM_UB:#map[0-9]*]] = affine_map<()[s0, s1] -> (s0, s1, 1024)>8 9// UJAM-FOUR-DAG: [[$UBMAP:#map[0-9]*]] = affine_map<()[s0] -> (s0 + 8)>10// UJAM-FOUR-DAG: [[$MAP_PLUS_1:#map[0-9]*]] = affine_map<(d0) -> (d0 + 1)>11// UJAM-FOUR-DAG: [[$MAP_PLUS_2:#map[0-9]*]] = affine_map<(d0) -> (d0 + 2)>12// UJAM-FOUR-DAG: [[$MAP_PLUS_3:#map[0-9]*]] = affine_map<(d0) -> (d0 + 3)>13 14// GPU-HJAM-DAG: [[$MAP_PLUS_1:#map[0-9]*]] = affine_map<(d0) -> (d0 + 1)>15 16// CHECK-LABEL: func @unroll_jam_imperfect_nest() {17func.func @unroll_jam_imperfect_nest() {18 affine.for %i = 0 to 101 {19 %x = "addi32"(%i, %i) : (index, index) -> i3220 affine.for %j = 0 to 17 {21 %y = "addi32"(%i, %i) : (index, index) -> i3222 %z = "addi32"(%y, %y) : (i32, i32) -> i3223 }24 %w = "foo"(%i, %x) : (index, i32) -> i3225 }26 return27}28// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 {29// CHECK-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]])30// CHECK-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])31// CHECK-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]])32// CHECK-NEXT: affine.for %{{.*}} = 0 to 17 {33// CHECK-NEXT: [[RES3:%[0-9]+]] = "addi32"([[IV0]], [[IV0]])34// CHECK-NEXT: "addi32"([[RES3]], [[RES3]]) : (i32, i32) -> i3235// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])36// CHECK-NEXT: [[RES4:%[0-9]+]] = "addi32"([[INC1]], [[INC1]])37// CHECK-NEXT: "addi32"([[RES4]], [[RES4]]) : (i32, i32) -> i3238// CHECK-NEXT: }39// CHECK: "foo"([[IV0]], [[RES1]])40// CHECK-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]])41// CHECK-NEXT: "foo"({{.*}}, [[RES2]])42// CHECK: }43// Cleanup loop (single iteration).44// CHECK: "addi32"(%c100, %c100)45// CHECK-NEXT: affine.for [[IV0]] = 0 to 17 {46// CHECK-NEXT: [[RESC:%[0-9]+]] = "addi32"(%c100, %c100)47// CHECK-NEXT: "addi32"([[RESC]], [[RESC]]) : (i32, i32) -> i3248// CHECK-NEXT: }49// CHECK-NEXT: "foo"(%c100, %{{.*}})50// CHECK-NEXT: return51 52gpu.module @unroll_jam {53 // GPU-HJAM-LABEL: func @unroll_jam_imperfect_nest() {54 gpu.func @unroll_jam_imperfect_nest() {55 affine.for %i = 0 to 101 {56 %x = "addi32"(%i, %i) : (index, index) -> i3257 affine.for %j = 0 to 17 {58 %y = "addi32"(%i, %i) : (index, index) -> i3259 %z = "addi32"(%y, %y) : (i32, i32) -> i3260 }61 %w = "foo"(%i, %x) : (index, i32) -> i3262 }63 gpu.return64 }65 // GPU-HJAM: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 {66 // GPU-HJAM-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]])67 // GPU-HJAM-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])68 // GPU-HJAM-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]])69 // GPU-HJAM-NEXT: affine.for %{{.*}} = 0 to 17 {70 // GPU-HJAM-NEXT: [[RES3:%[0-9]+]] = "addi32"([[IV0]], [[IV0]])71 // GPU-HJAM-NEXT: "addi32"([[RES3]], [[RES3]]) : (i32, i32) -> i3272 // GPU-HJAM-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])73 // GPU-HJAM-NEXT: [[RES4:%[0-9]+]] = "addi32"([[INC1]], [[INC1]])74 // GPU-HJAM-NEXT: "addi32"([[RES4]], [[RES4]]) : (i32, i32) -> i3275 // GPU-HJAM-NEXT: }76 // GPU-HJAM: "foo"([[IV0]], [[RES1]])77 // GPU-HJAM-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]])78 // GPU-HJAM-NEXT: "foo"({{.*}}, [[RES2]])79 // GPU-HJAM: }80 // Cleanup loop (single iteration).81 // GPU-HJAM: "addi32"(%c100, %c100)82 // GPU-HJAM-NEXT: affine.for [[IV0]] = 0 to 17 {83 // GPU-HJAM-NEXT: [[RESC:%[0-9]+]] = "addi32"(%c100, %c100)84 // GPU-HJAM-NEXT: "addi32"([[RESC]], [[RESC]]) : (i32, i32) -> i3285 // GPU-HJAM-NEXT: }86 // GPU-HJAM-NEXT: "foo"(%c100, %{{.*}})87 // GPU-HJAM-NEXT: return88}89 90// CHECK-LABEL: func @loop_nest_unknown_count_191// CHECK-SAME: [[N:arg[0-9]+]]: index92func.func @loop_nest_unknown_count_1(%N : index) {93 // CHECK-NEXT: affine.for %{{.*}} = 1 to [[$MAP_DIV_OFFSET]]()[%[[N]]] step 2 {94 // CHECK-NEXT: affine.for %{{.*}} = 1 to 100 {95 // CHECK-NEXT: "foo"() : () -> i3296 // CHECK-NEXT: "foo"() : () -> i3297 // CHECK-NEXT: }98 // CHECK-NEXT: }99 // A cleanup loop should be generated here.100 // CHECK-NEXT: affine.for %{{.*}} = [[$MAP_DIV_OFFSET]]()[%[[N]]] to %[[N]] {101 // CHECK-NEXT: affine.for %{{.*}} = 1 to 100 {102 // CHECK-NEXT: "foo"() : () -> i32103 // CHECK-NEXT: }104 // CHECK-NEXT: }105 affine.for %i = 1 to %N {106 affine.for %j = 1 to 100 {107 %x = "foo"() : () -> i32108 }109 }110 return111}112 113// UJAM-FOUR-LABEL: func @loop_nest_unknown_count_2114// UJAM-FOUR-SAME: %[[N:arg[0-9]+]]: index115func.func @loop_nest_unknown_count_2(%N : index) {116 // UJAM-FOUR-NEXT: affine.for [[IV0:%arg[0-9]+]] = %[[N]] to [[$UBMAP]]()[%[[N]]] step 4 {117 // UJAM-FOUR-NEXT: affine.for [[IV1:%arg[0-9]+]] = 1 to 100 {118 // UJAM-FOUR-NEXT: "foo"([[IV0]])119 // UJAM-FOUR-NEXT: [[IV_PLUS_1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])120 // UJAM-FOUR-NEXT: "foo"([[IV_PLUS_1]])121 // UJAM-FOUR-NEXT: [[IV_PLUS_2:%[0-9]+]] = affine.apply [[$MAP_PLUS_2]]([[IV0]])122 // UJAM-FOUR-NEXT: "foo"([[IV_PLUS_2]])123 // UJAM-FOUR-NEXT: [[IV_PLUS_3:%[0-9]+]] = affine.apply [[$MAP_PLUS_3]]([[IV0]])124 // UJAM-FOUR-NEXT: "foo"([[IV_PLUS_3]])125 // UJAM-FOUR-NEXT: }126 // UJAM-FOUR-NEXT: }127 // The cleanup loop is a single iteration one and is promoted.128 // UJAM-FOUR-NEXT: [[RES:%[0-9]+]] = affine.apply [[$UBMAP]]()[%[[N]]]129 // UJAM-FOUR-NEXT: affine.for [[IV0]] = 1 to 100 {130 // UJAM-FOUR-NEXT: "foo"([[RES]])131 // UJAM-FOUR-NEXT: }132 affine.for %i = %N to affine_map<()[s0] -> (s0+9)> ()[%N] {133 affine.for %j = 1 to 100 {134 "foo"(%i) : (index) -> ()135 }136 }137 return138}139 140// CHECK-LABEL: func @loop_nest_symbolic_and_min_upper_bound141// CHECK-SAME: [[M:arg[0-9]+]]: index142// CHECK-SAME: [[N:arg[0-9]+]]: index143// CHECK-SAME: [[K:arg[0-9]+]]: index144func.func @loop_nest_symbolic_and_min_upper_bound(%M : index, %N : index, %K : index) {145 affine.for %i = 0 to min affine_map<()[s0, s1] -> (s0, s1, 1024)>()[%M, %N] {146 affine.for %j = 0 to %K {147 "test.foo"(%i, %j) : (index, index) -> ()148 }149 }150 return151}152// No unroll-and-jam possible here as the lower bound for the cleanup loop won't153// be representable.154// CHECK-NEXT: affine.for [[IV0:%arg[0-9]+]] = 0 to min #map{{.*}}()[%[[M]], %[[N]]] {155// CHECK-NEXT: affine.for [[IV1:%arg[0-9]+]] = 0 to %[[K]] {156// CHECK-NEXT: "test.foo"([[IV0]], [[IV1]])157// CHECK-NEXT: }158// CHECK-NEXT: }159// CHECK-NEXT: return160 161// The inner loop trip count changes each iteration of outer loop.162// Do no unroll-and-jam.163// CHECK-LABEL: func @no_unroll_jam_dependent_ubound164func.func @no_unroll_jam_dependent_ubound(%in0: memref<?xf32, 1>) {165 affine.for %i = 0 to 100 {166 affine.for %k = 0 to affine_map<(d0) -> (d0 + 1)>(%i) {167 %y = "addi32"(%k, %k) : (index, index) -> i32168 }169 }170 return171}172// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 {173// CHECK-NEXT: affine.for [[IV1:%arg[0-9]+]] = 0 to [[$MAP_PLUS_1]]([[IV0]]) {174// CHECK-NEXT: "addi32"([[IV1]], [[IV1]])175// CHECK-NEXT: }176// CHECK-NEXT: }177// CHECK-NEXT: return178 179// Inner loop with one iter_arg.180// CHECK-LABEL: func @unroll_jam_one_iter_arg181func.func @unroll_jam_one_iter_arg() {182 affine.for %i = 0 to 101 {183 %cst = arith.constant 1 : i32184 %x = "addi32"(%i, %i) : (index, index) -> i32185 %red = affine.for %j = 0 to 17 iter_args(%acc = %cst) -> (i32) {186 %y = "bar"(%i, %j, %acc) : (index, index, i32) -> i32187 affine.yield %y : i32188 }189 %w = "foo"(%i, %x, %red) : (index, i32, i32) -> i32190 }191 return192}193// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 {194// CHECK-NEXT: [[CONST1:%[a-zA-Z0-9_]*]] = arith.constant 1 : i32195// CHECK-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]])196// CHECK-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])197// CHECK-NEXT: [[CONST2:%[a-zA-Z0-9_]*]] = arith.constant 1 : i32198// CHECK-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]])199// CHECK-NEXT: [[RES3:%[0-9]+]]:2 = affine.for [[IV1:%arg[0-9]+]] = 0 to 17 iter_args([[ACC1:%arg[0-9]+]] = [[CONST1]], [[ACC2:%arg[0-9]+]] = [[CONST2]]) -> (i32, i32) {200// CHECK-NEXT: [[RES4:%[0-9]+]] = "bar"([[IV0]], [[IV1]], [[ACC1]])201// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])202// CHECK-NEXT: [[RES5:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[ACC2]])203// CHECK-NEXT: affine.yield [[RES4]], [[RES5]]204// CHECK-NEXT: }205// CHECK: "foo"([[IV0]], [[RES1]], [[RES3]]#0)206// CHECK-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]])207// CHECK-NEXT: "foo"({{.*}}, [[RES2]], [[RES3]]#1)208// CHECK: }209// Cleanup loop (single iteration).210// CHECK: constant 1 : i32211// CHECK-NEXT: "addi32"(%c100, %c100)212// CHECK-NEXT: [[RES6:%[0-9]+]] = affine.for213// CHECK-NEXT: [[RES7:%[0-9]+]] = "bar"(%c100, {{.*}}, {{.*}})214// CHECK-NEXT: affine.yield [[RES7]] : i32215// CHECK-NEXT: }216// CHECK-NEXT: "foo"(%c100, %{{.*}}, [[RES6]])217// CHECK-NEXT: return218 219// Inner loop with multiple iter_args.220// CHECK-LABEL: func @unroll_jam_iter_args221func.func @unroll_jam_iter_args() {222 affine.for %i = 0 to 101 {223 %cst = arith.constant 0 : i32224 %cst1 = arith.constant 1 : i32225 %x = "addi32"(%i, %i) : (index, index) -> i32226 %red:2 = affine.for %j = 0 to 17 iter_args(%acc = %cst, %acc1 = %cst1) -> (i32, i32) {227 %y = "bar"(%i, %j, %acc) : (index, index, i32) -> i32228 %z = "bar1"(%i, %j, %acc1) : (index, index, i32) -> i32229 affine.yield %y, %z : i32, i32230 }231 %w = "foo"(%i, %x, %red#0, %red#1) : (index, i32, i32, i32) -> i32232 }233 return234}235// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 {236// CHECK-NEXT: [[CONST0:%[a-zA-Z0-9_]*]] = arith.constant 0 : i32237// CHECK-NEXT: [[CONST1:%[a-zA-Z0-9_]*]] = arith.constant 1 : i32238// CHECK-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]])239// CHECK-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])240// CHECK-NEXT: [[CONST2:%[a-zA-Z0-9_]*]] = arith.constant 0 : i32241// CHECK-NEXT: [[CONST3:%[a-zA-Z0-9_]*]] = arith.constant 1 : i32242// CHECK-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]])243// CHECK-NEXT: [[RES3:%[0-9]+]]:4 = affine.for [[IV1:%arg[0-9]+]] = 0 to 17 iter_args([[ACC0:%arg[0-9]+]] = [[CONST0]], [[ACC1:%arg[0-9]+]] = [[CONST1]],244// CHECK-SAME: [[ACC2:%arg[0-9]+]] = [[CONST2]], [[ACC3:%arg[0-9]+]] = [[CONST3]]) -> (i32, i32, i32, i32) {245// CHECK-NEXT: [[RES4:%[0-9]+]] = "bar"([[IV0]], [[IV1]], [[ACC0]])246// CHECK-NEXT: [[RES5:%[0-9]+]] = "bar1"([[IV0]], [[IV1]], [[ACC1]])247// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])248// CHECK-NEXT: [[RES6:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[ACC2]])249// CHECK-NEXT: [[RES7:%[0-9]+]] = "bar1"([[INC1]], [[IV1]], [[ACC3]])250// CHECK-NEXT: affine.yield [[RES4]], [[RES5]], [[RES6]], [[RES7]]251// CHECK-NEXT: }252// CHECK: "foo"([[IV0]], [[RES1]], [[RES3]]#0, [[RES3]]#1)253// CHECK-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]])254// CHECK-NEXT: "foo"({{.*}}, [[RES2]], [[RES3]]#2, [[RES3]]#3)255// CHECK: }256// Cleanup loop (single iteration).257// CHECK: constant 0 : i32258// CHECK-NEXT: constant 1 : i32259// CHECK-NEXT: "addi32"(%c100, %c100)260// CHECK-NEXT: [[RES8:%[0-9]+]]:2 = affine.for261// CHECK-NEXT: [[RES9:%[0-9]+]] = "bar"(%c100, {{.*}}, {{.*}})262// CHECK-NEXT: [[RES10:%[0-9]+]] = "bar1"(%c100, {{.*}}, {{.*}})263// CHECK-NEXT: affine.yield [[RES9]], [[RES10]] : i32, i32264// CHECK-NEXT: }265// CHECK-NEXT: "foo"(%c100, %{{.*}}, [[RES8]]#0, [[RES8]]#1)266// CHECK-NEXT: return267 268// When an iter operand is a function argument, do not replace any use of the269// operand .270// CHECK-LABEL: func @unroll_jam_iter_args_func_arg271// CHECK-SAME: [[INIT:%arg[0-9]+]]: i32272func.func @unroll_jam_iter_args_func_arg(%in: i32) {273 affine.for %i = 0 to 101 {274 %x = "addi32"(%i, %i) : (index, index) -> i32275 %red = affine.for %j = 0 to 17 iter_args(%acc = %in) -> (i32) {276 %y = "bar"(%i, %j, %acc) : (index, index, i32) -> i32277 affine.yield %y : i32278 }279 %w = "foo"(%i, %x, %red) : (index, i32, i32) -> i32280 }281 return282}283// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 {284// CHECK-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]])285// CHECK-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])286// CHECK-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]])287// CHECK-NEXT: [[RES3:%[0-9]+]]:2 = affine.for [[IV1:%arg[0-9]+]] = 0 to 17 iter_args([[ACC1:%arg[0-9]+]] = [[INIT]], [[ACC2:%arg[0-9]+]] = [[INIT]]) -> (i32, i32) {288// CHECK-NEXT: [[RES4:%[0-9]+]] = "bar"([[IV0]], [[IV1]], [[ACC1]])289// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])290// CHECK-NEXT: [[RES5:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[ACC2]])291// CHECK-NEXT: affine.yield [[RES4]], [[RES5]]292// CHECK-NEXT: }293// CHECK: "foo"([[IV0]], [[RES1]], [[RES3]]#0)294// CHECK-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]])295// CHECK-NEXT: "foo"({{.*}}, [[RES2]], [[RES3]]#1)296// CHECK: }297// Cleanup loop (single iteration).298// CHECK: "addi32"(%c100, %c100)299// CHECK-NEXT: [[RES6:%[0-9]+]] = affine.for300// CHECK-NEXT: [[RES7:%[0-9]+]] = "bar"(%c100, {{.*}}, {{.*}})301// CHECK-NEXT: affine.yield [[RES7]] : i32302// CHECK-NEXT: }303// CHECK-NEXT: "foo"(%c100, %{{.*}}, [[RES6]])304// CHECK-NEXT: return305 306// Nested inner loops, each with one iter_arg. The inner most loop uses its307// outer loop's iter_arg as its iter operand.308// CHECK-LABEL: func @unroll_jam_iter_args_nested309func.func @unroll_jam_iter_args_nested() {310 affine.for %i = 0 to 101 {311 %cst = arith.constant 1 : i32312 %x = "addi32"(%i, %i) : (index, index) -> i32313 %red = affine.for %j = 0 to 17 iter_args(%acc = %cst) -> (i32) {314 %red1 = affine.for %k = 0 to 35 iter_args(%acc1 = %acc) -> (i32) {315 %y = "bar"(%i, %j, %k, %acc1) : (index, index, index, i32) -> i32316 affine.yield %y : i32317 }318 affine.yield %red1 : i32319 }320 %w = "foo"(%i, %x, %red) : (index, i32, i32) -> i32321 }322 return323}324// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 {325// CHECK-NEXT: [[CONST1:%[a-zA-Z0-9_]*]] = arith.constant 1 : i32326// CHECK-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]])327// CHECK-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])328// CHECK-NEXT: [[CONST2:%[a-zA-Z0-9_]*]] = arith.constant 1 : i32329// CHECK-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]])330// CHECK-NEXT: [[RES3:%[0-9]+]]:2 = affine.for [[IV1:%arg[0-9]+]] = 0 to 17 iter_args([[ACC1:%arg[0-9]+]] = [[CONST1]], [[ACC2:%arg[0-9]+]] = [[CONST2]]) -> (i32, i32) {331// CHECK-NEXT: [[RES4:%[0-9]+]]:2 = affine.for [[IV2:%arg[0-9]+]] = 0 to 35 iter_args([[ACC3:%arg[0-9]+]] = [[ACC1]], [[ACC4:%arg[0-9]+]] = [[ACC2]]) -> (i32, i32) {332// CHECK-NEXT: [[RES5:%[0-9]+]] = "bar"([[IV0]], [[IV1]], [[IV2]], [[ACC3]])333// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])334// CHECK-NEXT: [[RES6:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[IV2]], [[ACC4]])335// CHECK-NEXT: affine.yield [[RES5]], [[RES6]]336// CHECK-NEXT: }337// CHECK-NEXT: affine.yield [[RES4]]#0, [[RES4]]#1338// CHECK-NEXT: }339// CHECK: "foo"([[IV0]], [[RES1]], [[RES3]]#0)340// CHECK-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]])341// CHECK-NEXT: "foo"({{.*}}, [[RES2]], [[RES3]]#1)342// CHECK: }343// Cleanup loop (single iteration).344// CHECK: constant 1 : i32345// CHECK-NEXT: "addi32"(%c100, %c100)346// CHECK-NEXT: [[RES6:%[0-9]+]] = affine.for347// CHECK-NEXT: [[RES7:%[0-9]+]] = affine.for348// CHECK-NEXT: [[RES8:%[0-9]+]] = "bar"(%c100, {{.*}}, {{.*}}, {{.*}})349// CHECK-NEXT: affine.yield [[RES8]] : i32350// CHECK-NEXT: }351// CHECK-NEXT: affine.yield [[RES7]] : i32352// CHECK-NEXT: }353// CHECK-NEXT: "foo"(%c100, %{{.*}}, [[RES6]])354// CHECK-NEXT: return355 356// Nested inner loops, each with one iter_arg. One loop uses its sibling loop's357// result as its iter operand.358// CHECK-LABEL: func @unroll_jam_iter_args_nested_affine_for_result359func.func @unroll_jam_iter_args_nested_affine_for_result() {360 affine.for %i = 0 to 101 {361 %cst = arith.constant 1 : i32362 %x = "addi32"(%i, %i) : (index, index) -> i32363 %red = affine.for %j = 0 to 17 iter_args(%acc = %cst) -> (i32) {364 %red1 = affine.for %k = 0 to 35 iter_args(%acc1 = %acc) -> (i32) {365 %y = "bar"(%i, %j, %k, %acc1) : (index, index, index, i32) -> i32366 affine.yield %acc : i32367 }368 %red2 = affine.for %l = 0 to 36 iter_args(%acc2 = %red1) -> (i32) {369 %y = "bar"(%i, %j, %l, %acc2) : (index, index, index, i32) -> i32370 affine.yield %y : i32371 }372 affine.yield %red2 : i32373 }374 %w = "foo"(%i, %x, %red) : (index, i32, i32) -> i32375 }376 return377}378// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 {379// CHECK-NEXT: [[CONST1:%[a-zA-Z0-9_]*]] = arith.constant 1 : i32380// CHECK-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]])381// CHECK-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])382// CHECK-NEXT: [[CONST2:%[a-zA-Z0-9_]*]] = arith.constant 1 : i32383// CHECK-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]])384// CHECK-NEXT: [[RES3:%[0-9]+]]:2 = affine.for [[IV1:%arg[0-9]+]] = 0 to 17 iter_args([[ACC1:%arg[0-9]+]] = [[CONST1]], [[ACC2:%arg[0-9]+]] = [[CONST2]]) -> (i32, i32) {385// CHECK-NEXT: [[RES4:%[0-9]+]]:2 = affine.for [[IV2:%arg[0-9]+]] = 0 to 35 iter_args([[ACC3:%arg[0-9]+]] = [[ACC1]], [[ACC4:%arg[0-9]+]] = [[ACC2]]) -> (i32, i32) {386// CHECK-NEXT: [[RES5:%[0-9]+]] = "bar"([[IV0]], [[IV1]], [[IV2]], [[ACC3]])387// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])388// CHECK-NEXT: [[RES6:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[IV2]], [[ACC4]])389// CHECK-NEXT: affine.yield [[ACC1]], [[ACC2]]390// CHECK-NEXT: }391// CHECK-NEXT: [[RES14:%[0-9]+]]:2 = affine.for [[IV3:%arg[0-9]+]] = 0 to 36 iter_args([[ACC13:%arg[0-9]+]] = [[RES4]]#0, [[ACC14:%arg[0-9]+]] = [[RES4]]#1) -> (i32, i32) {392// CHECK-NEXT: [[RES15:%[0-9]+]] = "bar"([[IV0]], [[IV1]], [[IV3]], [[ACC13]])393// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])394// CHECK-NEXT: [[RES16:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[IV3]], [[ACC14]])395// CHECK-NEXT: affine.yield [[RES15]], [[RES16]]396// CHECK-NEXT: }397// CHECK-NEXT: affine.yield [[RES14]]#0, [[RES14]]#1398// CHECK-NEXT: }399// CHECK: "foo"([[IV0]], [[RES1]], [[RES3]]#0)400// CHECK-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]])401// CHECK-NEXT: "foo"({{.*}}, [[RES2]], [[RES3]]#1)402// CHECK: }403// Cleanup loop (single iteration).404// CHECK: constant 1 : i32405// CHECK-NEXT: "addi32"(%c100, %c100)406// CHECK-NEXT: [[RES6:%[0-9]+]] = affine.for407// CHECK-NEXT: [[RES7:%[0-9]+]] = affine.for408// CHECK-NEXT: [[RES8:%[0-9]+]] = "bar"(%c100, {{.*}}, {{.*}}, {{.*}})409// CHECK-NEXT: affine.yield410// CHECK-NEXT: }411// CHECK-NEXT: [[RES17:%[0-9]+]] = affine.for412// CHECK-NEXT: [[RES18:%[0-9]+]] = "bar"(%c100, {{.*}}, {{.*}}, {{.*}})413// CHECK-NEXT: affine.yield [[RES18]] : i32414// CHECK-NEXT: }415// CHECK-NEXT: affine.yield [[RES17]] : i32416// CHECK-NEXT: }417// CHECK-NEXT: "foo"(%c100, %{{.*}}, [[RES6]])418// CHECK-NEXT: return419 420// Nested inner loops, each with one or more iter_args. Yeild the same value421// multiple times.422// CHECK-LABEL: func @unroll_jam_iter_args_nested_yield423func.func @unroll_jam_iter_args_nested_yield() {424 affine.for %i = 0 to 101 {425 %cst = arith.constant 1 : i32426 %x = "addi32"(%i, %i) : (index, index) -> i32427 %red:3 = affine.for %j = 0 to 17 iter_args(%acc = %cst, %acc1 = %cst, %acc2 = %cst) -> (i32, i32, i32) {428 %red1 = affine.for %k = 0 to 35 iter_args(%acc3 = %acc) -> (i32) {429 %y = "bar"(%i, %j, %k, %acc3) : (index, index, index, i32) -> i32430 affine.yield %y : i32431 }432 %red2:2 = affine.for %l = 0 to 36 iter_args(%acc4 = %acc1, %acc5 = %acc2) -> (i32, i32) {433 %y = "bar1"(%i, %j, %l, %acc4, %acc5) : (index, index, index, i32, i32) -> i32434 affine.yield %y, %y : i32, i32435 }436 affine.yield %red1, %red1, %red2#1 : i32, i32, i32437 }438 %w = "foo"(%i, %x, %red#0, %red#2) : (index, i32, i32, i32) -> i32439 }440 return441}442// CHECK: affine.for [[IV0:%arg[0-9]+]] = 0 to 100 step 2 {443// CHECK-NEXT: [[CONST1:%[a-zA-Z0-9_]*]] = arith.constant 1 : i32444// CHECK-NEXT: [[RES1:%[0-9]+]] = "addi32"([[IV0]], [[IV0]])445// CHECK-NEXT: [[INC:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])446// CHECK-NEXT: [[CONST2:%[a-zA-Z0-9_]*]] = arith.constant 1 : i32447// CHECK-NEXT: [[RES2:%[0-9]+]] = "addi32"([[INC]], [[INC]])448// CHECK-NEXT: [[RES3:%[0-9]+]]:6 = affine.for [[IV1:%arg[0-9]+]] = 0 to 17 iter_args([[ACC1:%arg[0-9]+]] = [[CONST1]], [[ACC2:%arg[0-9]+]] = [[CONST1]],449// CHECK-SAME: [[ACC3:%arg[0-9]+]] = [[CONST1]], [[ACC4:%arg[0-9]+]] = [[CONST2]], [[ACC5:%arg[0-9]+]] = [[CONST2]], [[ACC6:%arg[0-9]+]] = [[CONST2]]) -> (i32, i32, i32, i32, i32, i32) {450// CHECK-NEXT: [[RES4:%[0-9]+]]:2 = affine.for [[IV2:%arg[0-9]+]] = 0 to 35 iter_args([[ACC7:%arg[0-9]+]] = [[ACC1]], [[ACC8:%arg[0-9]+]] = [[ACC4]]) -> (i32, i32) {451// CHECK-NEXT: [[RES5:%[0-9]+]] = "bar"([[IV0]], [[IV1]], [[IV2]], [[ACC7]])452// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])453// CHECK-NEXT: [[RES6:%[0-9]+]] = "bar"([[INC1]], [[IV1]], [[IV2]], [[ACC8]])454// CHECK-NEXT: affine.yield [[RES5]], [[RES6]]455// CHECK-NEXT: }456// CHECK-NEXT: [[RES14:%[0-9]+]]:4 = affine.for [[IV3:%arg[0-9]+]] = 0 to 36 iter_args([[ACC13:%arg[0-9]+]] = [[ACC2]], [[ACC14:%arg[0-9]+]] = [[ACC3]],457// CHECK-SAME: [[ACC15:%arg[0-9]+]] = [[ACC5]], [[ACC16:%arg[0-9]+]] = [[ACC6]]) -> (i32, i32, i32, i32) {458// CHECK-NEXT: [[RES15:%[0-9]+]] = "bar1"([[IV0]], [[IV1]], [[IV3]], [[ACC13]], [[ACC14]])459// CHECK-NEXT: [[INC1:%[0-9]+]] = affine.apply [[$MAP_PLUS_1]]([[IV0]])460// CHECK-NEXT: [[RES16:%[0-9]+]] = "bar1"([[INC1]], [[IV1]], [[IV3]], [[ACC15]], [[ACC16]])461// CHECK-NEXT: affine.yield [[RES15]], [[RES15]], [[RES16]], [[RES16]]462// CHECK-NEXT: }463// CHECK-NEXT: affine.yield [[RES4]]#0, [[RES4]]#0, [[RES14]]#1, [[RES4]]#1, [[RES4]]#1, [[RES14]]#3464// CHECK-NEXT: }465// CHECK: "foo"([[IV0]], [[RES1]], [[RES3]]#0, [[RES3]]#2)466// CHECK-NEXT: affine.apply [[$MAP_PLUS_1]]([[IV0]])467// CHECK-NEXT: "foo"({{.*}}, [[RES2]], [[RES3]]#3, [[RES3]]#5)468// CHECK: }469// Cleanup loop (single iteration).470// CHECK: constant 1 : i32471// CHECK-NEXT: "addi32"(%c100, %c100)472// CHECK-NEXT: [[RES6:%[0-9]+]]:3 = affine.for473// CHECK-NEXT: [[RES7:%[0-9]+]] = affine.for474// CHECK-NEXT: [[RES8:%[0-9]+]] = "bar"(%c100, {{.*}}, {{.*}}, {{.*}})475// CHECK-NEXT: affine.yield [[RES8]] : i32476// CHECK-NEXT: }477// CHECK-NEXT: [[RES17:%[0-9]+]]:2 = affine.for478// CHECK-NEXT: [[RES18:%[0-9]+]] = "bar1"(%c100, {{.*}}, {{.*}}, {{.*}}, {{.*}})479// CHECK-NEXT: affine.yield [[RES18]], [[RES18]] : i32, i32480// CHECK-NEXT: }481// CHECK-NEXT: affine.yield [[RES7]], [[RES7]], [[RES17]]#1 : i32, i32, i32482// CHECK-NEXT: }483// CHECK-NEXT: "foo"(%c100, %{{.*}}, [[RES6]]#0, [[RES6]]#2)484// CHECK-NEXT: return485 486// CHECK-LABEL: func @unroll_jam_nested_iter_args_mulf487// CHECK-SAME: [[INIT0:%arg[0-9]+]]: f32, [[INIT1:%arg[0-9]+]]: f32488func.func @unroll_jam_nested_iter_args_mulf(%arg0: memref<21x30xf32, 1>, %init : f32, %init1 : f32) {489 %0 = affine.for %arg3 = 0 to 21 iter_args(%arg4 = %init) -> (f32) {490 %1 = affine.for %arg5 = 0 to 30 iter_args(%arg6 = %init1) -> (f32) {491 %3 = affine.load %arg0[%arg3, %arg5] : memref<21x30xf32, 1>492 %4 = arith.addf %arg6, %3 : f32493 affine.yield %4 : f32494 }495 %2 = arith.mulf %arg4, %1 : f32496 affine.yield %2 : f32497 }498 return499}500 501// CHECK: %[[CONST0:[a-zA-Z0-9_]*]] = arith.constant 20 : index502// CHECK-NEXT: [[RES:%[0-9]+]]:2 = affine.for %[[IV0:arg[0-9]+]] = 0 to 20 step 2 iter_args([[ACC0:%arg[0-9]+]] = [[INIT0]], [[ACC1:%arg[0-9]+]] = [[INIT0]]) -> (f32, f32) {503// CHECK-NEXT: [[RES1:%[0-9]+]]:2 = affine.for %[[IV1:arg[0-9]+]] = 0 to 30 iter_args([[ACC2:%arg[0-9]+]] = [[INIT1]], [[ACC3:%arg[0-9]+]] = [[INIT1]]) -> (f32, f32) {504// CHECK-NEXT: [[LOAD1:%[0-9]+]] = affine.load {{.*}}[%[[IV0]], %[[IV1]]]505// CHECK-NEXT: [[ADD1:%[0-9]+]] = arith.addf [[ACC2]], [[LOAD1]] : f32506// CHECK-NEXT: %[[INC1:[0-9]+]] = affine.apply [[$MAP_PLUS_1]](%[[IV0]])507// CHECK-NEXT: [[LOAD2:%[0-9]+]] = affine.load {{.*}}[%[[INC1]], %[[IV1]]]508// CHECK-NEXT: [[ADD2:%[0-9]+]] = arith.addf [[ACC3]], [[LOAD2]] : f32509// CHECK-NEXT: affine.yield [[ADD1]], [[ADD2]]510// CHECK-NEXT: }511// CHECK-NEXT: [[MUL1:%[0-9]+]] = arith.mulf [[ACC0]], [[RES1]]#0 : f32512// CHECK-NEXT: affine.apply513// CHECK-NEXT: [[MUL2:%[0-9]+]] = arith.mulf [[ACC1]], [[RES1]]#1 : f32514// CHECK-NEXT: affine.yield [[MUL1]], [[MUL2]]515// CHECK-NEXT: }516// Reduction op.517// CHECK-NEXT: [[MUL3:%[0-9]+]] = arith.mulf [[RES]]#0, [[RES]]#1 : f32518// Cleanup loop (single iteration).519// CHECK-NEXT: [[RES2:%[0-9]+]] = affine.for %[[IV2:arg[0-9]+]] = 0 to 30 iter_args([[ACC4:%arg[0-9]+]] = [[INIT1]]) -> (f32) {520// CHECK-NEXT: [[LOAD3:%[0-9]+]] = affine.load {{.*}}[%[[CONST0]], %[[IV2]]]521// CHECK-NEXT: [[ADD3:%[0-9]+]] = arith.addf [[ACC4]], [[LOAD3]] : f32522// CHECK-NEXT: affine.yield [[ADD3]] : f32523// CHECK-NEXT: }524// CHECK-NEXT: [[MUL4:%[0-9]+]] = arith.mulf [[MUL3]], [[RES2]] : f32525// CHECK-NEXT: return526 527// CHECK-LABEL: func @unroll_jam_iter_args_addi528// CHECK-SAME: [[INIT0:%arg[0-9]+]]: i32529func.func @unroll_jam_iter_args_addi(%arg0: memref<21xi32, 1>, %init : i32) {530 %0 = affine.for %arg3 = 0 to 21 iter_args(%arg4 = %init) -> (i32) {531 %1 = affine.load %arg0[%arg3] : memref<21xi32, 1>532 %2 = arith.addi %arg4, %1 : i32533 affine.yield %2 : i32534 }535 return536}537 538// CHECK: %[[CONST0:[a-zA-Z0-9_]*]] = arith.constant 20 : index539// CHECK-NEXT: [[RES:%[0-9]+]]:2 = affine.for %[[IV0:arg[0-9]+]] = 0 to 20 step 2 iter_args([[ACC0:%arg[0-9]+]] = [[INIT0]], [[ACC1:%arg[0-9]+]] = [[INIT0]]) -> (i32, i32) {540// CHECK-NEXT: [[LOAD1:%[0-9]+]] = affine.load {{.*}}[%[[IV0]]]541// CHECK-NEXT: [[ADD1:%[0-9]+]] = arith.addi [[ACC0]], [[LOAD1]] : i32542// CHECK-NEXT: %[[INC1:[0-9]+]] = affine.apply [[$MAP_PLUS_1]](%[[IV0]])543// CHECK-NEXT: [[LOAD2:%[0-9]+]] = affine.load {{.*}}[%[[INC1]]]544// CHECK-NEXT: [[ADD2:%[0-9]+]] = arith.addi [[ACC1]], [[LOAD2]] : i32545// CHECK-NEXT: affine.yield [[ADD1]], [[ADD2]]546// CHECK-NEXT: }547// Reduction op.548// CHECK-NEXT: [[ADD3:%[0-9]+]] = arith.addi [[RES]]#0, [[RES]]#1 : i32549// Cleanup loop (single iteration).550// CHECK-NEXT: [[LOAD3:%[0-9]+]] = affine.load {{.*}}[%[[CONST0]]]551// CHECK-NEXT: [[ADD4:%[0-9]+]] = arith.addi [[ADD3]], [[LOAD3]] : i32552// CHECK-NEXT: return553