[Mlir-commits] [mlir] [mlir] Fix support for loop normalization with integer indices (PR #76566)

[Arseniy Obolenskiy]

https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/76566

>From 1e928218c135a5e5ede1557a68d60b49e13c8d83 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <gooddoog at student.su>
Date: Fri, 29 Dec 2023 21:32:05 +0800
Subject: [PATCH] [mlir] Fix support for loop normalization with integer
 indices

Choose correct type for updated loop boundaries after scf loop normalization,
do not force chosen type to IndexType
---
 mlir/lib/Dialect/SCF/Utils/Utils.cpp     |  7 ++++--
 mlir/test/Dialect/SCF/transform-ops.mlir | 29 ++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index e85825595e3c1e..a2043c647d49a3 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -502,9 +502,12 @@ static LoopParams normalizeLoop(OpBuilder &boundsBuilder,
 
   Value newLowerBound =
       isZeroBased ? lowerBound
-                  : boundsBuilder.create<arith::ConstantIndexOp>(loc, 0);
+                  : boundsBuilder.create<arith::ConstantOp>(
+                        loc, boundsBuilder.getZeroAttr(lowerBound.getType()));
   Value newStep =
-      isStepOne ? step : boundsBuilder.create<arith::ConstantIndexOp>(loc, 1);
+      isStepOne ? step
+                : boundsBuilder.create<arith::ConstantOp>(
+                      loc, boundsBuilder.getIntegerAttr(step.getType(), 1));
 
   // Insert code computing the value of the original loop induction variable
   // from the "normalized" one.
diff --git a/mlir/test/Dialect/SCF/transform-ops.mlir b/mlir/test/Dialect/SCF/transform-ops.mlir
index 93ebf67f8b7133..0a626aa5a2b9ab 100644
--- a/mlir/test/Dialect/SCF/transform-ops.mlir
+++ b/mlir/test/Dialect/SCF/transform-ops.mlir
@@ -270,3 +270,32 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+// CHECK-LABEL: func @coalesce_i32_loops(
+
+func.func @coalesce_i32_loops() {
+  %0 = arith.constant 0 : i32
+  %1 = arith.constant 128 : i32
+  %2 = arith.constant 2 : i32
+  %3 = arith.constant 64 : i32
+  // CHECK: %[[C0_I32:.*]] = arith.constant 0 : i32
+  // CHECK: %[[C1_I32:.*]] = arith.constant 1 : i32
+  // CHECK: scf.for %[[ARG0:.*]] = %[[C0_I32]] to {{.*}} step %[[C1_I32]]  : i32
+  scf.for %i = %0 to %1 step %2 : i32 {
+    scf.for %j = %0 to %3 step %2 : i32 {
+      arith.addi %i, %j : i32
+    }
+  } {coalesce}
+  return
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["scf.for"]} attributes {coalesce} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %1 = transform.cast %0 : !transform.any_op to !transform.op<"scf.for">
+    %2 = transform.loop.coalesce %1: (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
+    transform.yield
+  }
+}