| // RUN: iree-dialects-opt --transform-dialect-interpreter %s | FileCheck %s |
| |
| |
| // CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> (s1 - (-s0 + s1) mod s2)> |
| // CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0] -> (-d0 + s0)> |
| // CHECK: func.func @fully_dynamic_bounds( |
| // CHECK-SAME: %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index |
| // CHECK: %[[C0_I32:.*]] = arith.constant 0 : i32 |
| // CHECK: %[[NEW_UB:.*]] = affine.apply #[[MAP0]]()[%[[LB]], %[[UB]], %[[STEP]]] |
| // CHECK: %[[CAST:.*]] = arith.index_cast %[[STEP]] : index to i32 |
| // CHECK: %[[LOOP:.*]] = scf.for %[[IV:.*]] = %[[LB]] to %[[NEW_UB]] |
| // CHECK-SAME: step %[[STEP]] iter_args(%[[ACC:.*]] = %[[C0_I32]]) -> (i32) { |
| // CHECK: %[[ADD:.*]] = arith.addi %[[ACC]], %[[CAST]] : i32 |
| // CHECK: scf.yield %[[ADD]] |
| // CHECK: } |
| // CHECK: %[[RESULT:.*]] = scf.for %[[IV2:.*]] = %[[NEW_UB]] to %[[UB]] |
| // CHECK-SAME: step %[[STEP]] iter_args(%[[ACC2:.*]] = %[[LOOP]]) -> (i32) { |
| // CHECK: %[[REM:.*]] = affine.apply #[[MAP1]](%[[IV2]])[%[[UB]]] |
| // CHECK: %[[CAST2:.*]] = arith.index_cast %[[REM]] |
| // CHECK: %[[ADD2:.*]] = arith.addi %[[ACC2]], %[[CAST2]] |
| // CHECK: scf.yield %[[ADD2]] |
| // CHECK: } |
| // CHECK: return %[[RESULT]] |
| #map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)> |
| func.func @fully_dynamic_bounds(%lb : index, %ub: index, %step: index) -> i32 { |
| %c0 = arith.constant 0 : i32 |
| %r = scf.for %iv = %lb to %ub step %step iter_args(%arg = %c0) -> i32 { |
| %s = affine.min #map(%ub, %iv)[%step] |
| %casted = arith.index_cast %s : index to i32 |
| %0 = arith.addi %arg, %casted : i32 |
| scf.yield %0 : i32 |
| } |
| return %r : i32 |
| } |
| |
| transform.with_pdl_patterns { |
| ^bb0(%arg0: !pdl.operation): |
| pdl.pattern @pdl_target : benefit(1) { |
| %args = operands |
| %results = types |
| %0 = pdl.operation "scf.for"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>) |
| %1 = pdl.attribute = @fully_dynamic_bounds |
| apply_native_constraint "nestedInFunc"(%0, %1 : !pdl.operation, !pdl.attribute) |
| // TODO: we don't want this, but it is the required terminator for pdl.pattern |
| rewrite %0 with "transform.dialect" |
| } |
| |
| transform.structured.canonicalized_sequence %arg0 { |
| ^bb1(%arg1: !pdl.operation): |
| %0 = pdl_match @pdl_target in %arg1 |
| transform.loop.peel %0 |
| } |
| } |
