521 lines · plain
1// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2' | FileCheck %s --check-prefix UNROLL-BY-22// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=3' | FileCheck %s --check-prefix UNROLL-BY-33// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=0' | FileCheck %s --check-prefix UNROLL-OUTER-BY-24// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=1' | FileCheck %s --check-prefix UNROLL-INNER-BY-25// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 annotate=true' | FileCheck %s --check-prefix UNROLL-BY-2-ANNOTATE6// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=6 unroll-up-to-factor=true}))" | FileCheck %s --check-prefix UNROLL-UP-TO7// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=5 cleanup-unroll=true}))" | FileCheck %s --check-prefix CLEANUP-UNROLL-BY-58// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll))" --split-input-file | FileCheck %s9 10func.func @dynamic_loop_unroll(%arg0 : index, %arg1 : index, %arg2 : index,11 %arg3: memref<?xf32>) {12 %0 = arith.constant 7.0 : f3213 scf.for %i0 = %arg0 to %arg1 step %arg2 {14 memref.store %0, %arg3[%i0] : memref<?xf32>15 }16 return17}18// UNROLL-BY-2-LABEL: func @dynamic_loop_unroll19// UNROLL-BY-2-SAME: %[[LB:.*0]]: index,20// UNROLL-BY-2-SAME: %[[UB:.*1]]: index,21// UNROLL-BY-2-SAME: %[[STEP:.*2]]: index,22// UNROLL-BY-2-SAME: %[[MEM:.*3]]: memref<?xf32>23//24// UNROLL-BY-2-DAG: %[[V0:.*]] = arith.subi %[[UB]], %[[LB]] : index25// UNROLL-BY-2-DAG: %[[C1:.*]] = arith.constant 1 : index26// UNROLL-BY-2-DAG: %[[V1:.*]] = arith.subi %[[STEP]], %[[C1]] : index27// UNROLL-BY-2-DAG: %[[V2:.*]] = arith.addi %[[V0]], %[[V1]] : index28// Compute trip count in V3.29// UNROLL-BY-2-DAG: %[[V3:.*]] = arith.divui %[[V2]], %[[STEP]] : index30// Store unroll factor in C2.31// UNROLL-BY-2-DAG: %[[C2:.*]] = arith.constant 2 : index32// UNROLL-BY-2-DAG: %[[V4:.*]] = arith.remsi %[[V3]], %[[C2]] : index33// UNROLL-BY-2-DAG: %[[V5:.*]] = arith.subi %[[V3]], %[[V4]] : index34// UNROLL-BY-2-DAG: %[[V6:.*]] = arith.muli %[[V5]], %[[STEP]] : index35// Compute upper bound of unrolled loop in V7.36// UNROLL-BY-2-DAG: %[[V7:.*]] = arith.addi %[[LB]], %[[V6]] : index37// Compute step of unrolled loop in V8.38// UNROLL-BY-2-DAG: %[[V8:.*]] = arith.muli %[[STEP]], %[[C2]] : index39// UNROLL-BY-2: scf.for %[[IV:.*]] = %[[LB]] to %[[V7]] step %[[V8]] {40// UNROLL-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>41// UNROLL-BY-2-NEXT: %[[C1_IV:.*]] = arith.constant 1 : index42// UNROLL-BY-2-NEXT: %[[V9:.*]] = arith.muli %[[STEP]], %[[C1_IV]] : index43// UNROLL-BY-2-NEXT: %[[V10:.*]] = arith.addi %[[IV]], %[[V9]] : index44// UNROLL-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V10]]] : memref<?xf32>45// UNROLL-BY-2-NEXT: }46// UNROLL-BY-2-NEXT: scf.for %[[IV:.*]] = %[[V7]] to %[[UB]] step %[[STEP]] {47// UNROLL-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>48// UNROLL-BY-2-NEXT: }49// UNROLL-BY-2-NEXT: return50 51// UNROLL-BY-3-LABEL: func @dynamic_loop_unroll52// UNROLL-BY-3-SAME: %[[LB:.*0]]: index,53// UNROLL-BY-3-SAME: %[[UB:.*1]]: index,54// UNROLL-BY-3-SAME: %[[STEP:.*2]]: index,55// UNROLL-BY-3-SAME: %[[MEM:.*3]]: memref<?xf32>56//57// UNROLL-BY-3-DAG: %[[V0:.*]] = arith.subi %[[UB]], %[[LB]] : index58// UNROLL-BY-3-DAG: %[[C1:.*]] = arith.constant 1 : index59// UNROLL-BY-3-DAG: %[[V1:.*]] = arith.subi %[[STEP]], %[[C1]] : index60// UNROLL-BY-3-DAG: %[[V2:.*]] = arith.addi %[[V0]], %[[V1]] : index61// Compute trip count in V3.62// UNROLL-BY-3-DAG: %[[V3:.*]] = arith.divui %[[V2]], %[[STEP]] : index63// Store unroll factor in C3.64// UNROLL-BY-3-DAG: %[[C3:.*]] = arith.constant 3 : index65// UNROLL-BY-3-DAG: %[[V4:.*]] = arith.remsi %[[V3]], %[[C3]] : index66// UNROLL-BY-3-DAG: %[[V5:.*]] = arith.subi %[[V3]], %[[V4]] : index67// UNROLL-BY-3-DAG: %[[V6:.*]] = arith.muli %[[V5]], %[[STEP]] : index68// Compute upper bound of unrolled loop in V7.69// UNROLL-BY-3-DAG: %[[V7:.*]] = arith.addi %[[LB]], %[[V6]] : index70// Compute step of unrolled loop in V8.71// UNROLL-BY-3-DAG: %[[V8:.*]] = arith.muli %[[STEP]], %[[C3]] : index72// UNROLL-BY-3: scf.for %[[IV:.*]] = %[[LB]] to %[[V7]] step %[[V8]] {73// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>74// UNROLL-BY-3-NEXT: %[[C1_IV:.*]] = arith.constant 1 : index75// UNROLL-BY-3-NEXT: %[[V9:.*]] = arith.muli %[[STEP]], %[[C1_IV]] : index76// UNROLL-BY-3-NEXT: %[[V10:.*]] = arith.addi %[[IV]], %[[V9]] : index77// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V10]]] : memref<?xf32>78// UNROLL-BY-3-NEXT: %[[C2_IV:.*]] = arith.constant 2 : index79// UNROLL-BY-3-NEXT: %[[V11:.*]] = arith.muli %[[STEP]], %[[C2_IV]] : index80// UNROLL-BY-3-NEXT: %[[V12:.*]] = arith.addi %[[IV]], %[[V11]] : index81// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V12]]] : memref<?xf32>82// UNROLL-BY-3-NEXT: }83// UNROLL-BY-3-NEXT: scf.for %[[IV:.*]] = %[[V7]] to %[[UB]] step %[[STEP]] {84// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>85// UNROLL-BY-3-NEXT: }86// UNROLL-BY-3-NEXT: return87 88func.func @dynamic_loop_unroll_outer_by_2(89 %arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index,90 %arg5 : index, %arg6: memref<?xf32>) {91 %0 = arith.constant 7.0 : f3292 scf.for %i0 = %arg0 to %arg1 step %arg2 {93 scf.for %i1 = %arg3 to %arg4 step %arg5 {94 memref.store %0, %arg6[%i1] : memref<?xf32>95 }96 }97 return98}99// UNROLL-OUTER-BY-2-LABEL: func @dynamic_loop_unroll_outer_by_2100// UNROLL-OUTER-BY-2-SAME: %[[LB0:.*0]]: index,101// UNROLL-OUTER-BY-2-SAME: %[[UB0:.*1]]: index,102// UNROLL-OUTER-BY-2-SAME: %[[STEP0:.*2]]: index,103// UNROLL-OUTER-BY-2-SAME: %[[LB1:.*3]]: index,104// UNROLL-OUTER-BY-2-SAME: %[[UB1:.*4]]: index,105// UNROLL-OUTER-BY-2-SAME: %[[STEP1:.*5]]: index,106// UNROLL-OUTER-BY-2-SAME: %[[MEM:.*6]]: memref<?xf32>107//108// UNROLL-OUTER-BY-2: scf.for %[[IV0:.*]] = %[[LB0]] to %{{.*}} step %{{.*}} {109// UNROLL-OUTER-BY-2-NEXT: scf.for %[[IV1:.*]] = %[[LB1]] to %[[UB1]] step %[[STEP1]] {110// UNROLL-OUTER-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV1]]] : memref<?xf32>111// UNROLL-OUTER-BY-2-NEXT: }112// UNROLL-OUTER-BY-2-NEXT: scf.for %[[IV1:.*]] = %[[LB1]] to %[[UB1]] step %[[STEP1]] {113// UNROLL-OUTER-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV1]]] : memref<?xf32>114// UNROLL-OUTER-BY-2-NEXT: }115// UNROLL-OUTER-BY-2-NEXT: }116// UNROLL-OUTER-BY-2-NEXT: scf.for %[[IV0:.*]] = %{{.*}} to %[[UB0]] step %[[STEP0]] {117// UNROLL-OUTER-BY-2-NEXT: scf.for %[[IV1:.*]] = %[[LB1]] to %[[UB1]] step %[[STEP1]] {118// UNROLL-OUTER-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV1]]] : memref<?xf32>119// UNROLL-OUTER-BY-2-NEXT: }120// UNROLL-OUTER-BY-2-NEXT: }121// UNROLL-OUTER-BY-2-NEXT: return122 123func.func @dynamic_loop_unroll_inner_by_2(124 %arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index,125 %arg5 : index, %arg6: memref<?xf32>) {126 %0 = arith.constant 7.0 : f32127 scf.for %i0 = %arg0 to %arg1 step %arg2 {128 scf.for %i1 = %arg3 to %arg4 step %arg5 {129 memref.store %0, %arg6[%i1] : memref<?xf32>130 }131 }132 return133}134// UNROLL-INNER-BY-2-LABEL: func @dynamic_loop_unroll_inner_by_2135// UNROLL-INNER-BY-2-SAME: %[[LB0:.*0]]: index,136// UNROLL-INNER-BY-2-SAME: %[[UB0:.*1]]: index,137// UNROLL-INNER-BY-2-SAME: %[[STEP0:.*2]]: index,138// UNROLL-INNER-BY-2-SAME: %[[LB1:.*3]]: index,139// UNROLL-INNER-BY-2-SAME: %[[UB1:.*4]]: index,140// UNROLL-INNER-BY-2-SAME: %[[STEP1:.*5]]: index,141// UNROLL-INNER-BY-2-SAME: %[[MEM:.*6]]: memref<?xf32>142//143// UNROLL-INNER-BY-2: scf.for %[[IV0:.*]] = %[[LB0]] to %[[UB0]] step %[[STEP0]] {144// UNROLL-INNER-BY-2: scf.for %[[IV1:.*]] = %[[LB1]] to %{{.*}} step %{{.*}} {145// UNROLL-INNER-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV1]]] : memref<?xf32>146// UNROLL-INNER-BY-2-NEXT: %[[C1_IV:.*]] = arith.constant 1 : index147// UNROLL-INNER-BY-2-NEXT: %[[V0:.*]] = arith.muli %[[STEP1]], %[[C1_IV]] : index148// UNROLL-INNER-BY-2-NEXT: %[[V1:.*]] = arith.addi %[[IV1]], %[[V0]] : index149// UNROLL-INNER-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>150// UNROLL-INNER-BY-2-NEXT: }151// UNROLL-INNER-BY-2-NEXT: scf.for %[[IV1:.*]] = %{{.*}} to %[[UB1]] step %[[STEP1]] {152// UNROLL-INNER-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV1]]] : memref<?xf32>153// UNROLL-INNER-BY-2-NEXT: }154// UNROLL-INNER-BY-2-NEXT: }155// UNROLL-INNER-BY-2-NEXT: return156 157// Test that no epilogue clean-up loop is generated because the trip count is158// a multiple of the unroll factor.159func.func @static_loop_unroll_by_2(%arg0 : memref<?xf32>) {160 %0 = arith.constant 7.0 : f32161 %lb = arith.constant 0 : index162 %ub = arith.constant 20 : index163 %step = arith.constant 1 : index164 scf.for %i0 = %lb to %ub step %step {165 memref.store %0, %arg0[%i0] : memref<?xf32>166 }167 return168}169// UNROLL-BY-2-LABEL: func @static_loop_unroll_by_2170// UNROLL-BY-2-SAME: %[[MEM:.*0]]: memref<?xf32>171//172// UNROLL-BY-2-DAG: %[[C0:.*]] = arith.constant 0 : index173// UNROLL-BY-2-DAG: %[[C1:.*]] = arith.constant 1 : index174// UNROLL-BY-2-DAG: %[[C20:.*]] = arith.constant 20 : index175// UNROLL-BY-2-DAG: %[[C2:.*]] = arith.constant 2 : index176// UNROLL-BY-2: scf.for %[[IV:.*]] = %[[C0]] to %[[C20]] step %[[C2]] {177// UNROLL-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>178// UNROLL-BY-2-NEXT: %[[C1_IV:.*]] = arith.constant 1 : index179// UNROLL-BY-2-NEXT: %[[V0:.*]] = arith.muli %[[C1]], %[[C1_IV]] : index180// UNROLL-BY-2-NEXT: %[[V1:.*]] = arith.addi %[[IV]], %[[V0]] : index181// UNROLL-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>182// UNROLL-BY-2-NEXT: }183// UNROLL-BY-2-NEXT: return184 185// UNROLL-BY-2-ANNOTATE-LABEL: func @static_loop_unroll_by_2186// UNROLL-BY-2-ANNOTATE: memref.store %{{.*}}, %[[MEM:.*0]][%{{.*}}] {unrolled_iteration = 0 : ui32} : memref<?xf32>187// UNROLL-BY-2-ANNOTATE: memref.store %{{.*}}, %[[MEM]][%{{.*}}] {unrolled_iteration = 1 : ui32} : memref<?xf32>188 189// Test that no epilogue clean-up loop is generated because the trip count190// (taking into account the non-unit step size) is a multiple of the unroll191// factor.192func.func @static_loop_step_2_unroll_by_2(%arg0 : memref<?xf32>) {193 %0 = arith.constant 7.0 : f32194 %lb = arith.constant 0 : index195 %ub = arith.constant 19 : index196 %step = arith.constant 2 : index197 scf.for %i0 = %lb to %ub step %step {198 memref.store %0, %arg0[%i0] : memref<?xf32>199 }200 return201}202 203// UNROLL-BY-2-LABEL: func @static_loop_step_2_unroll_by_2204// UNROLL-BY-2-SAME: %[[MEM:.*0]]: memref<?xf32>205//206// UNROLL-BY-2-DAG: %[[C0:.*]] = arith.constant 0 : index207// UNROLL-BY-2-DAG: %[[C2:.*]] = arith.constant 2 : index208// UNROLL-BY-2-DAG: %[[C19:.*]] = arith.constant 19 : index209// UNROLL-BY-2-DAG: %[[C4:.*]] = arith.constant 4 : index210// UNROLL-BY-2: scf.for %[[IV:.*]] = %[[C0]] to %[[C19]] step %[[C4]] {211// UNROLL-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>212// UNROLL-BY-2-NEXT: %[[C1_IV:.*]] = arith.constant 1 : index213// UNROLL-BY-2-NEXT: %[[V0:.*]] = arith.muli %[[C2]], %[[C1_IV]] : index214// UNROLL-BY-2-NEXT: %[[V1:.*]] = arith.addi %[[IV]], %[[V0]] : index215// UNROLL-BY-2-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>216// UNROLL-BY-2-NEXT: }217// UNROLL-BY-2-NEXT: return218 219// Test that epilogue clean up loop is generated (trip count is not220// a multiple of unroll factor).221func.func @static_loop_unroll_by_3(%arg0 : memref<?xf32>) {222 %0 = arith.constant 7.0 : f32223 %lb = arith.constant 0 : index224 %ub = arith.constant 20 : index225 %step = arith.constant 1 : index226 scf.for %i0 = %lb to %ub step %step {227 memref.store %0, %arg0[%i0] : memref<?xf32>228 }229 return230}231 232// UNROLL-BY-3-LABEL: func @static_loop_unroll_by_3233// UNROLL-BY-3-SAME: %[[MEM:.*0]]: memref<?xf32>234//235// UNROLL-BY-3-DAG: %[[C0:.*]] = arith.constant 0 : index236// UNROLL-BY-3-DAG: %[[C1:.*]] = arith.constant 1 : index237// UNROLL-BY-3-DAG: %[[C20:.*]] = arith.constant 20 : index238// UNROLL-BY-3-DAG: %[[C18:.*]] = arith.constant 18 : index239// UNROLL-BY-3-DAG: %[[C3:.*]] = arith.constant 3 : index240// UNROLL-BY-3: scf.for %[[IV:.*]] = %[[C0]] to %[[C18]] step %[[C3]] {241// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>242// UNROLL-BY-3-NEXT: %[[C1_IV:.*]] = arith.constant 1 : index243// UNROLL-BY-3-NEXT: %[[V0:.*]] = arith.muli %[[C1]], %[[C1_IV]] : index244// UNROLL-BY-3-NEXT: %[[V1:.*]] = arith.addi %[[IV]], %[[V0]] : index245// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>246// UNROLL-BY-3-NEXT: %[[C2_IV:.*]] = arith.constant 2 : index247// UNROLL-BY-3-NEXT: %[[V2:.*]] = arith.muli %[[C1]], %[[C2_IV]] : index248// UNROLL-BY-3-NEXT: %[[V3:.*]] = arith.addi %[[IV]], %[[V2]] : index249// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V3]]] : memref<?xf32>250// UNROLL-BY-3-NEXT: }251// UNROLL-BY-3-NEXT: scf.for %[[IV:.*]] = %[[C18]] to %[[C20]] step %[[C1]] {252// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>253// UNROLL-BY-3-NEXT: }254// UNROLL-BY-3-NEXT: return255 256// Test that the single iteration epilogue loop body is promoted to the loops257// containing block.258func.func @static_loop_unroll_by_3_promote_epilogue(%arg0 : memref<?xf32>) {259 %0 = arith.constant 7.0 : f32260 %lb = arith.constant 0 : index261 %ub = arith.constant 10 : index262 %step = arith.constant 1 : index263 scf.for %i0 = %lb to %ub step %step {264 memref.store %0, %arg0[%i0] : memref<?xf32>265 }266 return267}268// UNROLL-BY-3-LABEL: func @static_loop_unroll_by_3_promote_epilogue269// UNROLL-BY-3-SAME: %[[MEM:.*0]]: memref<?xf32>270//271// UNROLL-BY-3-DAG: %[[C0:.*]] = arith.constant 0 : index272// UNROLL-BY-3-DAG: %[[C1:.*]] = arith.constant 1 : index273// UNROLL-BY-3-DAG: %[[C10:.*]] = arith.constant 10 : index274// UNROLL-BY-3-DAG: %[[C9:.*]] = arith.constant 9 : index275// UNROLL-BY-3-DAG: %[[C3:.*]] = arith.constant 3 : index276// UNROLL-BY-3: scf.for %[[IV:.*]] = %[[C0]] to %[[C9]] step %[[C3]] {277// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[IV]]] : memref<?xf32>278// UNROLL-BY-3-NEXT: %[[C1_IV:.*]] = arith.constant 1 : index279// UNROLL-BY-3-NEXT: %[[V0:.*]] = arith.muli %[[C1]], %[[C1_IV]] : index280// UNROLL-BY-3-NEXT: %[[V1:.*]] = arith.addi %[[IV]], %[[V0]] : index281// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>282// UNROLL-BY-3-NEXT: %[[C2_IV:.*]] = arith.constant 2 : index283// UNROLL-BY-3-NEXT: %[[V2:.*]] = arith.muli %[[C1]], %[[C2_IV]] : index284// UNROLL-BY-3-NEXT: %[[V3:.*]] = arith.addi %[[IV]], %[[V2]] : index285// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V3]]] : memref<?xf32>286// UNROLL-BY-3-NEXT: }287// UNROLL-BY-3-NEXT: memref.store %{{.*}}, %[[MEM]][%[[C9]]] : memref<?xf32>288// UNROLL-BY-3-NEXT: return289 290// Test unroll-up-to functionality.291func.func @static_loop_unroll_up_to_factor(%arg0 : memref<?xf32>) {292 %0 = arith.constant 7.0 : f32293 %lb = arith.constant 0 : index294 %ub = arith.constant 2 : index295 affine.for %i0 = %lb to %ub {296 affine.store %0, %arg0[%i0] : memref<?xf32>297 }298 return299}300// UNROLL-UP-TO-LABEL: func @static_loop_unroll_up_to_factor301// UNROLL-UP-TO-SAME: %[[MEM:.*0]]: memref<?xf32>302//303// UNROLL-UP-TO-DAG: %[[C0:.*]] = arith.constant 0 : index304// UNROLL-UP-TO-DAG: %[[C2:.*]] = arith.constant 2 : index305// UNROLL-UP-TO-NEXT: %[[V0:.*]] = affine.apply {{.*}}306// UNROLL-UP-TO-NEXT: affine.store %{{.*}}, %[[MEM]][%[[V0]]] : memref<?xf32>307// UNROLL-UP-TO-NEXT: %[[V1:.*]] = affine.apply {{.*}}308// UNROLL-UP-TO-NEXT: affine.store %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>309// UNROLL-UP-TO-NEXT: return310 311// Test that epilogue's arguments are correctly renamed.312func.func @static_loop_unroll_by_3_rename_epilogue_arguments() -> (f32, f32) {313 %0 = arith.constant 7.0 : f32314 %lb = arith.constant 0 : index315 %ub = arith.constant 20 : index316 %step = arith.constant 1 : index317 %result:2 = scf.for %i0 = %lb to %ub step %step iter_args(%arg0 = %0, %arg1 = %0) -> (f32, f32) {318 %add = arith.addf %arg0, %arg1 : f32319 %mul = arith.mulf %arg0, %arg1 : f32320 scf.yield %add, %mul : f32, f32321 }322 return %result#0, %result#1 : f32, f32323}324// UNROLL-BY-3-LABEL: func @static_loop_unroll_by_3_rename_epilogue_arguments325//326// UNROLL-BY-3-DAG: %[[CST:.*]] = arith.constant {{.*}} : f32327// UNROLL-BY-3-DAG: %[[C0:.*]] = arith.constant 0 : index328// UNROLL-BY-3-DAG: %[[C1:.*]] = arith.constant 1 : index329// UNROLL-BY-3-DAG: %[[C20:.*]] = arith.constant 20 : index330// UNROLL-BY-3-DAG: %[[C18:.*]] = arith.constant 18 : index331// UNROLL-BY-3-DAG: %[[C3:.*]] = arith.constant 3 : index332// UNROLL-BY-3: %[[FOR:.*]]:2 = scf.for %[[IV:.*]] = %[[C0]] to %[[C18]] step %[[C3]]333// UNROLL-BY-3-SAME: iter_args(%[[ARG0:.*]] = %[[CST]], %[[ARG1:.*]] = %[[CST]]) -> (f32, f32) {334// UNROLL-BY-3-NEXT: %[[ADD0:.*]] = arith.addf %[[ARG0]], %[[ARG1]] : f32335// UNROLL-BY-3-NEXT: %[[MUL0:.*]] = arith.mulf %[[ARG0]], %[[ARG1]] : f32336// UNROLL-BY-3-NEXT: %[[ADD1:.*]] = arith.addf %[[ADD0]], %[[MUL0]] : f32337// UNROLL-BY-3-NEXT: %[[MUL1:.*]] = arith.mulf %[[ADD0]], %[[MUL0]] : f32338// UNROLL-BY-3-NEXT: %[[ADD2:.*]] = arith.addf %[[ADD1]], %[[MUL1]] : f32339// UNROLL-BY-3-NEXT: %[[MUL2:.*]] = arith.mulf %[[ADD1]], %[[MUL1]] : f32340// UNROLL-BY-3-NEXT: scf.yield %[[ADD2]], %[[MUL2]] : f32, f32341// UNROLL-BY-3-NEXT: }342// UNROLL-BY-3: %[[EFOR:.*]]:2 = scf.for %[[EIV:.*]] = %[[C18]] to %[[C20]] step %[[C1]]343// UNROLL-BY-3-SAME: iter_args(%[[EARG0:.*]] = %[[FOR]]#0, %[[EARG1:.*]] = %[[FOR]]#1) -> (f32, f32) {344// UNROLL-BY-3-NEXT: %[[EADD:.*]] = arith.addf %[[EARG0]], %[[EARG1]] : f32345// UNROLL-BY-3-NEXT: %[[EMUL:.*]] = arith.mulf %[[EARG0]], %[[EARG1]] : f32346// UNROLL-BY-3-NEXT: scf.yield %[[EADD]], %[[EMUL]] : f32, f32347// UNROLL-BY-3-NEXT: }348// UNROLL-BY-3-NEXT: return %[[EFOR]]#0, %[[EFOR]]#1 : f32, f32349 350// Test that epilogue clean up loop is generated (trip count is less351// than an unroll factor).352func.func @static_loop_unroll_by_5_with_cleanup(%arg0 : memref<?xf32>) {353 %0 = arith.constant 7.0 : f32354 %lb = arith.constant 0 : index355 %ub = arith.constant 3 : index356 affine.for %i0 = %lb to %ub {357 memref.store %0, %arg0[%i0] : memref<?xf32>358 }359 return360}361 362// CLEANUP-UNROLL-BY-5-LABEL: func @static_loop_unroll_by_5_with_cleanup363// CLEANUP-UNROLL-BY-5-SAME: %[[MEM:.*0]]: memref<?xf32>364//365// CLEANUP-UNROLL-BY-5-DAG: %[[C0:.*]] = arith.constant 0 : index366// CLEANUP-UNROLL-BY-5-DAG: %[[C3:.*]] = arith.constant 3 : index367// CLEANUP-UNROLL-BY-5-NEXT: %[[V0:.*]] = affine.apply {{.*}}368// CLEANUP-UNROLL-BY-5-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V0]]] : memref<?xf32>369// CLEANUP-UNROLL-BY-5-NEXT: %[[V1:.*]] = affine.apply {{.*}}370// CLEANUP-UNROLL-BY-5-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>371// CLEANUP-UNROLL-BY-5-NEXT: %[[V2:.*]] = affine.apply {{.*}}372// CLEANUP-UNROLL-BY-5-NEXT: memref.store %{{.*}}, %[[MEM]][%[[V2]]] : memref<?xf32>373// CLEANUP-UNROLL-BY-5-NEXT: return374 375// -----376 377// Test loop unrolling when the yielded value remains unchanged.378// CHECK: [[$MAP:#map]] = affine_map<(d0) -> (-d0 + 64, (d0 floordiv 8) ceildiv 64, -d0 - 16, d0 * -64)>379// CHECK-LABEL: func @loop_unroll_static_yield_value380func.func @loop_unroll_static_yield_value_test1() {381 %true_4 = arith.constant true382 %c1 = arith.constant 1 : index383 %103 = affine.for %arg2 = 0 to 40 iter_args(%arg3 = %true_4) -> (i1) {384 %324 = affine.max affine_map<(d0) -> (-d0 + 64, (d0 floordiv 8) ceildiv 64, -d0 - 16, d0 * -64)>(%c1)385 affine.yield %true_4 : i1386 }387 return388}389// CHECK: %[[TRUE:.*]] = arith.constant true390// CHECK-NEXT: %[[C1:.*]] = arith.constant 1 : index391// CHECK-NEXT: affine.for %{{.*}} = 0 to 40 step 4 iter_args(%{{.*}} = %[[TRUE]]) -> (i1) {392// CHECK-NEXT: affine.max [[$MAP]](%[[C1]])393// CHECK-NEXT: affine.max [[$MAP]](%[[C1]])394// CHECK-NEXT: affine.max [[$MAP]](%[[C1]])395// CHECK-NEXT: affine.max [[$MAP]](%[[C1]])396// CHECK-NEXT: affine.yield %[[TRUE]] : i1397// CHECK-NEXT: }398// CHECK-NEXT: return399 400// -----401 402// Loop unrolling when the yielded value is loop iv.403// CHECK: [[$MAP0:#map[0-9]*]] = affine_map<(d0) -> (-d0 + 64, (d0 floordiv 8) ceildiv 64, -d0 - 16, d0 * -64)>404// CHECK: [[$MAP1:#map[0-9]*]] = affine_map<(d0) -> (d0 + 2)>405// CHECK: [[$MAP2:#map[0-9]*]] = affine_map<(d0) -> (d0 + 4)>406// CHECK: [[$MAP3:#map[0-9]*]] = affine_map<(d0) -> (d0 + 6)>407// CHECK-LABEL: func @loop_unroll_yield_loop_iv408func.func @loop_unroll_yield_loop_iv() {409 %c1 = arith.constant 1 : index410 %103 = affine.for %arg2 = 0 to 40 step 2 iter_args(%arg3 = %c1) -> (index) {411 %324 = affine.max affine_map<(d0) -> (-d0 + 64, (d0 floordiv 8) ceildiv 64, -d0 - 16, d0 * -64)>(%arg2)412 affine.yield %arg2 : index413 }414 return415}416// CHECK: %[[C1:.*]] = arith.constant 1 : index417// CHECK-NEXT: affine.for %[[LOOP_IV:.*]] = 0 to 40 step 8 iter_args(%{{.*}} = %[[C1]]) -> (index) {418// CHECK-NEXT: affine.max [[$MAP0]](%[[LOOP_IV]])419// CHECK-NEXT: %[[LOOP_IV_PLUS_2:.*]] = affine.apply [[$MAP1]](%[[LOOP_IV]])420// CHECK-NEXT: affine.max [[$MAP0]](%[[LOOP_IV_PLUS_2]])421// CHECK-NEXT: %[[LOOP_IV_PLUS_4:.*]] = affine.apply [[$MAP2]](%[[LOOP_IV]])422// CHECK-NEXT: affine.max [[$MAP0]](%[[LOOP_IV_PLUS_4]])423// CHECK-NEXT: %[[LOOP_IV_PLUS_6:.*]] = affine.apply [[$MAP3]](%[[LOOP_IV]])424// CHECK-NEXT: affine.max [[$MAP0]](%[[LOOP_IV_PLUS_6]])425// CHECK-NEXT: affine.yield %[[LOOP_IV]] : index426// CHECK-NEXT: }427// CHECK-NEXT: return428 429// -----430 431// Loop unrolling when the yielded value is iter_arg.432// CHECK: [[$MAP:#map]] = affine_map<(d0) -> (-d0 + 64, (d0 floordiv 8) ceildiv 64, -d0 - 16, d0 * -64)>433// CHECK-LABEL: func @loop_unroll_yield_iter_arg434func.func @loop_unroll_yield_iter_arg() {435 %c1 = arith.constant 1 : index436 %103 = affine.for %arg2 = 0 to 40 step 2 iter_args(%arg3 = %c1) -> (index) {437 %324 = affine.max affine_map<(d0) -> (-d0 + 64, (d0 floordiv 8) ceildiv 64, -d0 - 16, d0 * -64)>(%arg3)438 affine.yield %arg3 : index439 }440 return441}442// CHECK: %[[C1:.*]] = arith.constant 1 : index443// CHECK-NEXT: affine.for %{{.*}} = 0 to 40 step 8 iter_args(%[[ITER_ARG:.*]] = %[[C1]]) -> (index) {444// CHECK-NEXT: affine.max [[$MAP]](%[[ITER_ARG]])445// CHECK-NEXT: affine.max [[$MAP]](%[[ITER_ARG]])446// CHECK-NEXT: affine.max [[$MAP]](%[[ITER_ARG]])447// CHECK-NEXT: affine.max [[$MAP]](%[[ITER_ARG]])448// CHECK-NEXT: affine.yield %[[ITER_ARG]] : index449// CHECK-NEXT: }450// CHECK-NEXT: return451 452// -----453 454// Test the loop unroller works with integer IV type.455func.func @static_loop_unroll_with_integer_iv() -> (f32, f32) {456 %0 = arith.constant 7.0 : f32457 %lb = arith.constant 0 : i32458 %ub = arith.constant 20 : i32459 %step = arith.constant 1 : i32460 %result:2 = scf.for %i0 = %lb to %ub step %step iter_args(%arg0 = %0, %arg1 = %0) -> (f32, f32) : i32{461 %add = arith.addf %arg0, %arg1 : f32462 %mul = arith.mulf %arg0, %arg1 : f32463 scf.yield %add, %mul : f32, f32464 }465 return %result#0, %result#1 : f32, f32466}467// UNROLL-BY-3-LABEL: func @static_loop_unroll_with_integer_iv468//469// UNROLL-BY-3-DAG: %[[CST:.*]] = arith.constant {{.*}} : f32470// UNROLL-BY-3-DAG: %[[C0:.*]] = arith.constant 0 : i32471// UNROLL-BY-3-DAG: %[[C1:.*]] = arith.constant 1 : i32472// UNROLL-BY-3-DAG: %[[C20:.*]] = arith.constant 20 : i32473// UNROLL-BY-3-DAG: %[[C18:.*]] = arith.constant 18 : i32474// UNROLL-BY-3-DAG: %[[C3:.*]] = arith.constant 3 : i32475// UNROLL-BY-3: %[[FOR:.*]]:2 = scf.for %[[IV:.*]] = %[[C0]] to %[[C18]] step %[[C3]]476// UNROLL-BY-3-SAME: iter_args(%[[ARG0:.*]] = %[[CST]], %[[ARG1:.*]] = %[[CST]]) -> (f32, f32) : i32 {477// UNROLL-BY-3-NEXT: %[[ADD0:.*]] = arith.addf %[[ARG0]], %[[ARG1]] : f32478// UNROLL-BY-3-NEXT: %[[MUL0:.*]] = arith.mulf %[[ARG0]], %[[ARG1]] : f32479// UNROLL-BY-3-NEXT: %[[ADD1:.*]] = arith.addf %[[ADD0]], %[[MUL0]] : f32480// UNROLL-BY-3-NEXT: %[[MUL1:.*]] = arith.mulf %[[ADD0]], %[[MUL0]] : f32481// UNROLL-BY-3-NEXT: %[[ADD2:.*]] = arith.addf %[[ADD1]], %[[MUL1]] : f32482// UNROLL-BY-3-NEXT: %[[MUL2:.*]] = arith.mulf %[[ADD1]], %[[MUL1]] : f32483// UNROLL-BY-3-NEXT: scf.yield %[[ADD2]], %[[MUL2]] : f32, f32484// UNROLL-BY-3-NEXT: }485// UNROLL-BY-3: %[[EFOR:.*]]:2 = scf.for %[[EIV:.*]] = %[[C18]] to %[[C20]] step %[[C1]]486// UNROLL-BY-3-SAME: iter_args(%[[EARG0:.*]] = %[[FOR]]#0, %[[EARG1:.*]] = %[[FOR]]#1) -> (f32, f32) : i32 {487// UNROLL-BY-3-NEXT: %[[EADD:.*]] = arith.addf %[[EARG0]], %[[EARG1]] : f32488// UNROLL-BY-3-NEXT: %[[EMUL:.*]] = arith.mulf %[[EARG0]], %[[EARG1]] : f32489// UNROLL-BY-3-NEXT: scf.yield %[[EADD]], %[[EMUL]] : f32, f32490// UNROLL-BY-3-NEXT: }491// UNROLL-BY-3-NEXT: return %[[EFOR]]#0, %[[EFOR]]#1 : f32, f32492 493// -----494 495// Test loop unrolling when the yielded value is defined above the loop.496func.func @loop_unroll_static_yield_value_defined_above(%init: i32) {497 %c42 = arith.constant 42 : i32498 %c0 = arith.constant 0 : index499 %c1 = arith.constant 1 : index500 %c4 = arith.constant 4 : index501 %103:2 = scf.for %i = %c0 to %c4 step %c1502 iter_args(%iter1 = %c42, %iter2 = %init) -> (i32, i32) {503 %0 = arith.andi %iter2, %iter1 : i32504 scf.yield %0, %init : i32, i32505 }506 return507}508// UNROLL-OUTER-BY-2-LABEL: @loop_unroll_static_yield_value_defined_above(509// UNROLL-OUTER-BY-2-SAME: %[[INIT:.*]]: i32) {510// UNROLL-OUTER-BY-2-DAG: %[[C42:.*]] = arith.constant 42 : i32511// UNROLL-OUTER-BY-2-DAG: %[[C0:.*]] = arith.constant 0 : index512// UNROLL-OUTER-BY-2-DAG: %[[C4:.*]] = arith.constant 4 : index513// UNROLL-OUTER-BY-2-DAG: %[[C2:.*]] = arith.constant 2 : index514// UNROLL-OUTER-BY-2: scf.for %{{.*}} = %[[C0]] to %[[C4]] step %[[C2]]515// UNROLL-OUTER-BY-2-SAME: iter_args(%[[ITER1:.*]] = %[[C42]],516// UNROLL-OUTER-BY-2-SAME: %[[ITER2:.*]] = %[[INIT]])517// UNROLL-OUTER-BY-2: %[[SUM:.*]] = arith.andi %[[ITER2]], %[[ITER1]]518// UNROLL-OUTER-BY-2: %[[SUM1:.*]] = arith.andi %[[INIT]], %[[SUM]]519// UNROLL-OUTER-BY-2: scf.yield %[[SUM1]], %[[INIT]] : i32, i32520 521