[Mlir-commits] [mlir] [mlir][scf] Specialize loops with constant affine.min operands (PR #228909)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-scf

Author: 陈子昂 (Michael-Chen-NJU)

<details>
<summary>Changes</summary>

`scf-for-loop-specialization` only inspected affine map results, so it missed constants supplied through map operands. Canonicalize the map and operands before finding constant results.

Fixes #<!-- -->228470

---
Full diff: https://github.com/llvm/llvm-project/pull/228909.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp (+4-1) 
- (modified) mlir/test/Dialect/SCF/for-loop-specialization.mlir (+26) 


``````````diff
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index a0b5ae9b66091..685c57319a216 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -84,8 +84,11 @@ static void specializeForLoopForUnrolling(ForOp op) {
   auto minOp = bound.getDefiningOp<AffineMinOp>();
   if (!minOp)
     return;
+  AffineMap canonicalizedMap = minOp.getMap();
+  SmallVector<Value> operands(minOp.getMapOperands());
+  affine::canonicalizeMapAndOperands(&canonicalizedMap, &operands);
   int64_t minConstant = std::numeric_limits<int64_t>::max();
-  for (AffineExpr expr : minOp.getMap().getResults()) {
+  for (AffineExpr expr : canonicalizedMap.getResults()) {
     if (auto constantIndex = dyn_cast<AffineConstantExpr>(expr))
       minConstant = std::min(minConstant, constantIndex.getValue());
   }
diff --git a/mlir/test/Dialect/SCF/for-loop-specialization.mlir b/mlir/test/Dialect/SCF/for-loop-specialization.mlir
index ff66c6c1e47f5..59e522f621963 100644
--- a/mlir/test/Dialect/SCF/for-loop-specialization.mlir
+++ b/mlir/test/Dialect/SCF/for-loop-specialization.mlir
@@ -37,3 +37,29 @@ func.func @for(%outer: index, %A: memref<?xf32>, %B: memref<?xf32>,
 // CHECK:           }
 // CHECK:           return
 // CHECK:         }
+
+// -----
+
+func.func @for_constant_operand(%n: index) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c8 = arith.constant 8 : index
+  %ub = affine.min affine_map<(d0)[s0] -> (d0, s0)>(%n)[%c8]
+  scf.for %i = %c0 to %ub step %c1 {
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @for_constant_operand(
+// CHECK-SAME:    [[N:%.*]]: index)
+// CHECK:         [[C8:%.*]] = arith.constant 8 : index
+// CHECK:         [[MIN:%.*]] = affine.min #{{.*}}([[N]]){{\[}}[[C8]]]
+// CHECK:         [[SPECIALIZED_C8:%.*]] = arith.constant 8 : index
+// CHECK:         [[PRED:%.*]] = arith.cmpi eq, [[MIN]], [[SPECIALIZED_C8]]
+// CHECK:         scf.if [[PRED]] {
+// CHECK:           scf.for {{%.*}} = {{%.*}} to [[SPECIALIZED_C8]] step {{%.*}} {
+// CHECK:           }
+// CHECK:         } else {
+// CHECK:           scf.for {{%.*}} = {{%.*}} to [[MIN]] step {{%.*}} {
+// CHECK:           }
+// CHECK:         }

``````````

</details>


https://github.com/llvm/llvm-project/pull/228909