1// RUN: mlir-opt %s -scf-for-loop-peeling=peel-front=true -split-input-file | FileCheck %s 2 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 : index 6// CHECK-DAG: %[[C0_I32:.*]] = arith.constant 0 : i32 7// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 8// CHECK-DAG: %[[C17:.*]] = arith.constant 17 : index 9// 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]](%[[C4]], %[[IV]])[%[[C4]]] 12// CHECK: %[[CAST:.*]] = arith.index_cast %[[MIN]] : index to i32 13// CHECK: %[[INIT:.*]] = arith.addi %[[ACC]], %[[CAST]] : i32 14// 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 i32 20// CHECK: %[[ADD:.*]] = arith.addi %[[ACC2]], %[[CAST2]] : i32 21// 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 : i32 27 %lb = arith.constant 0 : index 28 %step = arith.constant 4 : index 29 %ub = arith.constant 17 : index 30 %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 i32 33 %0 = arith.addi %arg, %casted : i32 34 scf.yield %0 : i32 35 } 36 return %r : i32 37} 38 39// ----- 40 41// CHECK-DAG: #[[MAP:.*]] = affine_map<(d0, d1)[s0] -> (4, d0 - d1)> 42// CHECK: func @no_loop_results( 43// CHECK-SAME: %[[UB:.*]]: index, %[[MEMREF:.*]]: memref<i32> 44// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index 45// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 46// CHECK: scf.for %[[IV:.*]] = %[[C0]] to %[[C4]] step %[[C4]] { 47// CHECK: %[[MIN:.*]] = affine.min #[[MAP]](%[[C4]], %[[IV]])[%[[C4]]] 48// CHECK: %[[LOAD:.*]] = memref.load %[[MEMREF]][] 49// CHECK: %[[CAST:.*]] = arith.index_cast %[[MIN]] 50// CHECK: %[[ADD:.*]] = arith.addi %[[LOAD]], %[[CAST]] : i32 51// CHECK: memref.store %[[ADD]], %[[MEMREF]] 52// CHECK: } 53// CHECK: scf.for %[[IV2:.*]] = %[[C4]] to %[[UB]] step %[[C4]] { 54// CHECK: %[[REM:.*]] = affine.min #[[MAP]](%[[UB]], %[[IV2]])[%[[C4]]] 55// CHECK: %[[LOAD2:.*]] = memref.load %[[MEMREF]][] 56// CHECK: %[[CAST2:.*]] = arith.index_cast %[[REM]] 57// CHECK: %[[ADD2:.*]] = arith.addi %[[LOAD2]], %[[CAST2]] 58// CHECK: memref.store %[[ADD2]], %[[MEMREF]] 59// CHECK: } 60// CHECK: return 61#map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)> 62func.func @no_loop_results(%ub : index, %d : memref<i32>) { 63 %c0_i32 = arith.constant 0 : i32 64 %lb = arith.constant 0 : index 65 %step = arith.constant 4 : index 66 scf.for %iv = %lb to %ub step %step { 67 %s = affine.min #map(%ub, %iv)[%step] 68 %r = memref.load %d[] : memref<i32> 69 %casted = arith.index_cast %s : index to i32 70 %0 = arith.addi %r, %casted : i32 71 memref.store %0, %d[] : memref<i32> 72 } 73 return 74} 75 76// ----- 77 78// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0, s1] -> (s0 + s1)> 79// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)> 80// CHECK: func @fully_dynamic_bounds( 81// CHECK-SAME: %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index 82// CHECK: %[[C0_I32:.*]] = arith.constant 0 : i32 83// CHECK: %[[NEW_UB:.*]] = affine.apply #[[MAP0]]()[%[[LB]], %[[STEP]]] 84// CHECK: %[[FIRST:.*]] = scf.for %[[IV:.*]] = %[[LB]] to %[[NEW_UB]] 85// CHECK-SAME: step %[[STEP]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) { 86// CHECK: %[[MIN:.*]] = affine.min #[[MAP1]](%[[NEW_UB]], %[[IV]])[%[[STEP]]] 87// CHECK: %[[CAST:.*]] = arith.index_cast %[[MIN]] : index to i32 88// CHECK: %[[ADD:.*]] = arith.addi %[[ACC]], %[[CAST]] : i32 89// CHECK: scf.yield %[[ADD]] 90// CHECK: } 91// CHECK: %[[RESULT:.*]] = scf.for %[[IV2:.*]] = %[[NEW_UB]] to %[[UB]] 92// CHECK-SAME: step %[[STEP]] iter_args(%[[ACC2:.*]] = %[[FIRST]]) -> (i32) { 93// CHECK: %[[REM:.*]] = affine.min #[[MAP1]](%[[UB]], %[[IV2]])[%[[STEP]]] 94// CHECK: %[[CAST2:.*]] = arith.index_cast %[[REM]] 95// CHECK: %[[ADD2:.*]] = arith.addi %[[ACC2]], %[[CAST2]] 96// CHECK: scf.yield %[[ADD2]] 97// CHECK: } 98// CHECK: return %[[RESULT]] 99#map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)> 100func.func @fully_dynamic_bounds(%lb : index, %ub: index, %step: index) -> i32 { 101 %c0 = arith.constant 0 : i32 102 %r = scf.for %iv = %lb to %ub step %step iter_args(%arg = %c0) -> i32 { 103 %s = affine.min #map(%ub, %iv)[%step] 104 %casted = arith.index_cast %s : index to i32 105 %0 = arith.addi %arg, %casted : i32 106 scf.yield %0 : i32 107 } 108 return %r : i32 109} 110 111// ----- 112 113// CHECK-DAG: #[[MAP:.*]] = affine_map<(d0, d1)[s0] -> (4, d0 - d1)> 114// CHECK: func @two_iteration_example( 115// CHECK-DAG: %[[C0_I32:.*]] = arith.constant 0 : i32 116// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index 117// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index 118// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index 119// CHECK-DAG: %[[C6:.*]] = arith.constant 6 : index 120// CHECK: %[[FIRST:.*]] = scf.for %[[IV:.*]] = %[[C2]] to %[[C6]] 121// CHECK-SAME: step %[[C4]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) { 122// CHECK: %[[MIN:.*]] = affine.min #[[MAP]](%[[C6]], %[[IV]])[%[[C4]]] 123// CHECK: %[[CAST:.*]] = arith.index_cast %[[MIN]] : index to i32 124// CHECK: %[[INIT:.*]] = arith.addi %[[ACC]], %[[CAST]] : i32 125// CHECK: scf.yield %[[INIT]] 126// CHECK: } 127// CHECK: %[[RESULT:.*]] = scf.for %[[IV2:.*]] = %[[C6]] to %[[C8]] 128// CHECK-SAME: step %[[C4]] iter_args(%[[ACC2:.*]] = %[[FIRST]]) -> (i32) { 129// CHECK: %[[MIN2:.*]] = affine.min #[[MAP]](%[[C8]], %[[IV2]])[%[[C4]]] 130// CHECK: %[[CAST2:.*]] = arith.index_cast %[[MIN2]] : index to i32 131// CHECK: %[[ADD:.*]] = arith.addi %[[ACC2]], %[[CAST2]] : i32 132// CHECK: scf.yield %[[ADD]] 133// CHECK: } 134// CHECK: return %[[RESULT]] 135#map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)> 136func.func @two_iteration_example() -> i32 { 137 %c0_i32 = arith.constant 0 : i32 138 %lb = arith.constant 2 : index 139 %step = arith.constant 4 : index 140 %ub = arith.constant 8 : index 141 %r = scf.for %iv = %lb to %ub step %step iter_args(%arg = %c0_i32) -> i32 { 142 %s = affine.min #map(%ub, %iv)[%step] 143 %casted = arith.index_cast %s : index to i32 144 %0 = arith.addi %arg, %casted : i32 145 scf.yield %0 : i32 146 } 147 return %r : i32 148} 149 150// ----- 151 152// CHECK-DAG: #[[MAP:.*]] = affine_map<(d0, d1)[s0] -> (4, d0 - d1)> 153// CHECK: func @no_peeling_front( 154// CHECK-DAG: %[[C0_I32:.*]] = arith.constant 0 : i32 155// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 156// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index 157// CHECK: %[[RESULT:.*]] = scf.for %[[IV:.*]] = %[[C0]] to %[[C4]] 158// CHECK-SAME: step %[[C4]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) { 159// CHECK: %[[MIN:.*]] = affine.min #[[MAP]](%[[C4]], %[[IV]])[%[[C4]]] 160// CHECK: %[[CAST:.*]] = arith.index_cast %[[MIN]] : index to i32 161// CHECK: %[[ADD:.*]] = arith.addi %[[ACC]], %[[CAST]] : i32 162// CHECK: scf.yield %[[ADD]] 163// CHECK: } 164// CHECK: return %[[RESULT]] 165#map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)> 166func.func @no_peeling_front() -> i32 { 167 %c0_i32 = arith.constant 0 : i32 168 %lb = arith.constant 0 : index 169 %step = arith.constant 4 : index 170 %ub = arith.constant 4 : index 171 %r = scf.for %iv = %lb to %ub step %step iter_args(%arg = %c0_i32) -> i32 { 172 %s = affine.min #map(%ub, %iv)[%step] 173 %casted = arith.index_cast %s : index to i32 174 %0 = arith.addi %arg, %casted : i32 175 scf.yield %0 : i32 176 } 177 return %r : i32 178} 179