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

[llvmlistbot at llvm.org]

Author: 陈子昂
Date: 2026-10-04T23:27:22Z
New Revision: 9c8979d1bb551949cf7985bf36c53e55aaf57ee2

URL: https://github.com/llvm/llvm-project/commit/9c8979d1bb551949cf7985bf36c53e55aaf57ee2
DIFF: https://github.com/llvm/llvm-project/commit/9c8979d1bb551949cf7985bf36c53e55aaf57ee2.diff

LOG: [mlir][scf] Specialize loops with constant affine.min operands (#228909)

`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

Added: 
    

Modified: 
    mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
    mlir/test/Dialect/SCF/for-loop-specialization.mlir

Removed: 
    


################################################################################
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:         }