221 lines · plain
1// RUN: mlir-opt %s -scf-for-loop-peeling=peel-front=true -split-input-file | FileCheck %s2 3// CHECK-DAG: #[[MAP:.*]] = affine_map<(d0, d1)[s0] -> (4, d0 - d1)>4// CHECK: func @fully_static_bounds(5// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index6// CHECK-DAG: %[[C0_I32:.*]] = arith.constant 0 : i327// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index8// CHECK-DAG: %[[C17:.*]] = arith.constant 17 : index9// CHECK: %[[FIRST:.*]] = scf.for %[[IV:.*]] = %[[C0]] to %[[C4]]10// CHECK-SAME: step %[[C4]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) {11// CHECK: %[[MIN:.*]] = affine.min #[[MAP]](%[[C17]], %[[IV]])[%[[C4]]]12// CHECK: %[[CAST:.*]] = arith.index_cast %[[MIN]] : index to i3213// CHECK: %[[INIT:.*]] = arith.addi %[[ACC]], %[[CAST]] : i3214// CHECK: scf.yield %[[INIT]]15// CHECK: }16// CHECK: %[[RESULT:.*]] = scf.for %[[IV2:.*]] = %[[C4]] to %[[C17]]17// CHECK-SAME: step %[[C4]] iter_args(%[[ACC2:.*]] = %[[FIRST]]) -> (i32) {18// CHECK: %[[MIN2:.*]] = affine.min #[[MAP]](%[[C17]], %[[IV2]])[%[[C4]]]19// CHECK: %[[CAST2:.*]] = arith.index_cast %[[MIN2]] : index to i3220// CHECK: %[[ADD:.*]] = arith.addi %[[ACC2]], %[[CAST2]] : i3221// CHECK: scf.yield %[[ADD]]22// CHECK: }23// CHECK: return %[[RESULT]]24#map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>25func.func @fully_static_bounds() -> i32 {26 %c0_i32 = arith.constant 0 : i3227 %lb = arith.constant 0 : index28 %step = arith.constant 4 : index29 %ub = arith.constant 17 : index30 %r = scf.for %iv = %lb to %ub step %step iter_args(%arg = %c0_i32) -> i32 {31 %s = affine.min #map(%ub, %iv)[%step]32 %casted = arith.index_cast %s : index to i3233 %0 = arith.addi %arg, %casted : i3234 scf.yield %0 : i3235 }36 return %r : i3237}38// -----39 40// CHECK-LABEL: func.func @static_two_iterations_ub_used_in_loop(41// CHECK-SAME: %[[IFM1:.*]]: memref<1xi32>) -> i32 {42// CHECK: %[[C0_I32:.*]] = arith.constant 0 : i3243// CHECK: %[[C0_IDX:.*]] = arith.constant 0 : index44// CHECK: %[[C4_IDX:.*]] = arith.constant 4 : index45// CHECK: %[[C7_IDX:.*]] = arith.constant 7 : index46// CHECK: %[[FOR1:.*]] = scf.for %[[IV1:.*]] = %[[C0_IDX]] to %[[C4_IDX]] step %[[C4_IDX]] iter_args(%[[ARG1:.*]] = %[[C0_I32]]) -> (i32) {47// CHECK: %[[AFFINE_MIN1:.*]] = affine.min #map(%[[C7_IDX]], %[[IV1]]){{\[}}%[[C4_IDX]]]48// CHECK: %[[CAST1:.*]] = arith.index_cast %[[AFFINE_MIN1]] : index to i3249// CHECK: %[[ADDI1:.*]] = arith.addi %[[ARG1]], %[[CAST1]] : i3250// CHECK: %[[LOAD1:.*]] = memref.load %[[IFM1]]{{\[}}%[[C7_IDX]]] : memref<1xi32>51// CHECK: %[[ADDI2:.*]] = arith.addi %[[ADDI1]], %[[LOAD1]] : i3252// CHECK: scf.yield %[[ADDI2]] : i3253// CHECK: }54// CHECK: %[[FOR2:.*]] = scf.for %[[IV2:.*]] = %[[C4_IDX]] to %[[C7_IDX]] step %[[C4_IDX]] iter_args(%[[ARG2:.*]] = %[[RESULT1:.*]]) -> (i32) {55// CHECK: %[[AFFINE_MIN2:.*]] = affine.min #map(%[[C7_IDX]], %[[IV2]]){{\[}}%[[C4_IDX]]]56// CHECK: %[[CAST2:.*]] = arith.index_cast %[[AFFINE_MIN2]] : index to i3257// CHECK: %[[ADDI3:.*]] = arith.addi %[[ARG2]], %[[CAST2]] : i3258// CHECK: %[[LOAD2:.*]] = memref.load %[[IFM1]]{{\[}}%[[C7_IDX]]] : memref<1xi32>59// CHECK: %[[ADDI4:.*]] = arith.addi %[[ADDI3]], %[[LOAD2]] : i3260// CHECK: scf.yield %[[ADDI4]] : i3261// CHECK: }62// CHECK: return %[[RESULT2:.*]] : i3263// CHECK: }64 65#map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>66func.func @static_two_iterations_ub_used_in_loop(%arg0: memref<1xi32>) -> i32 {67 %c0_i32 = arith.constant 0 : i3268 %lb = arith.constant 0 : index69 %step = arith.constant 4 : index70 %ub = arith.constant 7 : index71 %r = scf.for %iv = %lb to %ub step %step iter_args(%arg = %c0_i32) -> i32 {72 %s = affine.min #map(%ub, %iv)[%step]73 %casted = arith.index_cast %s : index to i3274 %0 = arith.addi %arg, %casted : i3275 %1 = memref.load %arg0[%ub] : memref<1xi32>76 %result = arith.addi %0, %1 : i3277 scf.yield %result : i3278 }79 return %r : i3280}81// -----82 83// CHECK-DAG: #[[MAP:.*]] = affine_map<(d0, d1)[s0] -> (4, d0 - d1)>84// CHECK: func @no_loop_results(85// CHECK-SAME: %[[UB:.*]]: index, %[[MEMREF:.*]]: memref<i32>86// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index87// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index88// CHECK: scf.for %[[IV:.*]] = %[[C0]] to %[[C4]] step %[[C4]] {89// CHECK: %[[MIN:.*]] = affine.min #[[MAP]](%[[UB]], %[[IV]])[%[[C4]]]90// CHECK: %[[LOAD:.*]] = memref.load %[[MEMREF]][]91// CHECK: %[[CAST:.*]] = arith.index_cast %[[MIN]]92// CHECK: %[[ADD:.*]] = arith.addi %[[LOAD]], %[[CAST]] : i3293// CHECK: memref.store %[[ADD]], %[[MEMREF]]94// CHECK: }95// CHECK: scf.for %[[IV2:.*]] = %[[C4]] to %[[UB]] step %[[C4]] {96// CHECK: %[[REM:.*]] = affine.min #[[MAP]](%[[UB]], %[[IV2]])[%[[C4]]]97// CHECK: %[[LOAD2:.*]] = memref.load %[[MEMREF]][]98// CHECK: %[[CAST2:.*]] = arith.index_cast %[[REM]]99// CHECK: %[[ADD2:.*]] = arith.addi %[[LOAD2]], %[[CAST2]]100// CHECK: memref.store %[[ADD2]], %[[MEMREF]]101// CHECK: }102// CHECK: return103#map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>104func.func @no_loop_results(%ub : index, %d : memref<i32>) {105 %c0_i32 = arith.constant 0 : i32106 %lb = arith.constant 0 : index107 %step = arith.constant 4 : index108 scf.for %iv = %lb to %ub step %step {109 %s = affine.min #map(%ub, %iv)[%step]110 %r = memref.load %d[] : memref<i32>111 %casted = arith.index_cast %s : index to i32112 %0 = arith.addi %r, %casted : i32113 memref.store %0, %d[] : memref<i32>114 }115 return116}117 118// -----119 120// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0, s1] -> (s0 + s1)>121// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>122// CHECK: func @fully_dynamic_bounds(123// CHECK-SAME: %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index124// CHECK: %[[C0_I32:.*]] = arith.constant 0 : i32125// CHECK: %[[NEW_UB:.*]] = affine.apply #[[MAP0]]()[%[[LB]], %[[STEP]]]126// CHECK: %[[FIRST:.*]] = scf.for %[[IV:.*]] = %[[LB]] to %[[NEW_UB]]127// CHECK-SAME: step %[[STEP]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) {128// CHECK: %[[MIN:.*]] = affine.min #[[MAP1]](%[[UB]], %[[IV]])[%[[STEP]]]129// CHECK: %[[CAST:.*]] = arith.index_cast %[[MIN]] : index to i32130// CHECK: %[[ADD:.*]] = arith.addi %[[ACC]], %[[CAST]] : i32131// CHECK: scf.yield %[[ADD]]132// CHECK: }133// CHECK: %[[RESULT:.*]] = scf.for %[[IV2:.*]] = %[[NEW_UB]] to %[[UB]]134// CHECK-SAME: step %[[STEP]] iter_args(%[[ACC2:.*]] = %[[FIRST]]) -> (i32) {135// CHECK: %[[REM:.*]] = affine.min #[[MAP1]](%[[UB]], %[[IV2]])[%[[STEP]]]136// CHECK: %[[CAST2:.*]] = arith.index_cast %[[REM]]137// CHECK: %[[ADD2:.*]] = arith.addi %[[ACC2]], %[[CAST2]]138// CHECK: scf.yield %[[ADD2]]139// CHECK: }140// CHECK: return %[[RESULT]]141#map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>142func.func @fully_dynamic_bounds(%lb : index, %ub: index, %step: index) -> i32 {143 %c0 = arith.constant 0 : i32144 %r = scf.for %iv = %lb to %ub step %step iter_args(%arg = %c0) -> i32 {145 %s = affine.min #map(%ub, %iv)[%step]146 %casted = arith.index_cast %s : index to i32147 %0 = arith.addi %arg, %casted : i32148 scf.yield %0 : i32149 }150 return %r : i32151}152 153// -----154 155// CHECK-DAG: #[[MAP:.*]] = affine_map<(d0, d1)[s0] -> (4, d0 - d1)>156// CHECK: func @two_iteration_example(157// CHECK-DAG: %[[C0_I32:.*]] = arith.constant 0 : i32158// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index159// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index160// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index161// CHECK-DAG: %[[C6:.*]] = arith.constant 6 : index162// CHECK: %[[FIRST:.*]] = scf.for %[[IV:.*]] = %[[C2]] to %[[C6]]163// CHECK-SAME: step %[[C4]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) {164// CHECK: %[[MIN:.*]] = affine.min #[[MAP]](%[[C8]], %[[IV]])[%[[C4]]]165// CHECK: %[[CAST:.*]] = arith.index_cast %[[MIN]] : index to i32166// CHECK: %[[INIT:.*]] = arith.addi %[[ACC]], %[[CAST]] : i32167// CHECK: scf.yield %[[INIT]]168// CHECK: }169// CHECK: %[[RESULT:.*]] = scf.for %[[IV2:.*]] = %[[C6]] to %[[C8]]170// CHECK-SAME: step %[[C4]] iter_args(%[[ACC2:.*]] = %[[FIRST]]) -> (i32) {171// CHECK: %[[MIN2:.*]] = affine.min #[[MAP]](%[[C8]], %[[IV2]])[%[[C4]]]172// CHECK: %[[CAST2:.*]] = arith.index_cast %[[MIN2]] : index to i32173// CHECK: %[[ADD:.*]] = arith.addi %[[ACC2]], %[[CAST2]] : i32174// CHECK: scf.yield %[[ADD]]175// CHECK: }176// CHECK: return %[[RESULT]]177#map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>178func.func @two_iteration_example() -> i32 {179 %c0_i32 = arith.constant 0 : i32180 %lb = arith.constant 2 : index181 %step = arith.constant 4 : index182 %ub = arith.constant 8 : index183 %r = scf.for %iv = %lb to %ub step %step iter_args(%arg = %c0_i32) -> i32 {184 %s = affine.min #map(%ub, %iv)[%step]185 %casted = arith.index_cast %s : index to i32186 %0 = arith.addi %arg, %casted : i32187 scf.yield %0 : i32188 }189 return %r : i32190}191 192// -----193 194// CHECK-DAG: #[[MAP:.*]] = affine_map<(d0, d1)[s0] -> (4, d0 - d1)>195// CHECK: func @no_peeling_front(196// CHECK-DAG: %[[C0_I32:.*]] = arith.constant 0 : i32197// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index198// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index199// CHECK: %[[RESULT:.*]] = scf.for %[[IV:.*]] = %[[C0]] to %[[C4]]200// CHECK-SAME: step %[[C4]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) {201// CHECK: %[[MIN:.*]] = affine.min #[[MAP]](%[[C4]], %[[IV]])[%[[C4]]]202// CHECK: %[[CAST:.*]] = arith.index_cast %[[MIN]] : index to i32203// CHECK: %[[ADD:.*]] = arith.addi %[[ACC]], %[[CAST]] : i32204// CHECK: scf.yield %[[ADD]]205// CHECK: }206// CHECK: return %[[RESULT]]207#map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>208func.func @no_peeling_front() -> i32 {209 %c0_i32 = arith.constant 0 : i32210 %lb = arith.constant 0 : index211 %step = arith.constant 4 : index212 %ub = arith.constant 4 : index213 %r = scf.for %iv = %lb to %ub step %step iter_args(%arg = %c0_i32) -> i32 {214 %s = affine.min #map(%ub, %iv)[%step]215 %casted = arith.index_cast %s : index to i32216 %0 = arith.addi %arg, %casted : i32217 scf.yield %0 : i32218 }219 return %r : i32220}221