[Mlir-commits] [mlir] [mlir] computeSliceParameters: Fix offset when m(0) != 0 (PR #122492)

[Matthias Gehre]

https://github.com/mgehre-amd updated https://github.com/llvm/llvm-project/pull/122492

>From 86b891d14704a27a70feff8befad04f6a5a5548d Mon Sep 17 00:00:00 2001
From: Matthias Gehre <matthias.gehre at amd.com>
Date: Fri, 10 Jan 2025 16:49:38 +0100
Subject: [PATCH 1/3] For affine maps where `m(0) != 0`, like `affine_map<(d0)
 -> (d0 + 3)` in ```   %generic = linalg.generic     {indexing_maps =
 [affine_map<(d0) -> (d0 + 3)>,                       affine_map<(d0) ->
 (d0)>],      iterator_types = ["parallel"]} ins(%arg0: tensor<9xf32>)
 outs(%empty : tensor<6xf32>) {     ^bb0(%in : f32, %out: f32):      
 linalg.yield %in : f32     } -> tensor<6xf32> ``` tiling currently computes
 the wrong slice offsets. When tiling above example with a size of 3, it would
 compute ``` scf.for %i = ...   %slice = tensor.extract_slice %arg0[%i + 3]
 [6] [1]   linalg.generic     {indexing_maps = [affine_map<(d0) -> (d0 + 3)>, 
                      affine_map<(d0) -> (d0)>],      iterator_types =
 ["parallel"]} ins(%slice: tensor<6xf32>) ``` and thus apply the `+3` twice
 (once in the extract slice and a second time in the linalg.generic).

This PR fixes this to yield an offset of
`tensor.extract_slice %arg0[%i] [6] [1]` instead.
---
 mlir/lib/Dialect/Linalg/Utils/Utils.cpp   | 11 +++++++++-
 mlir/test/Dialect/Linalg/tile-offset.mlir | 25 +++++++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/Linalg/tile-offset.mlir

diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
index 38e427af1c4846..2ebeccb1505048 100644
--- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
@@ -596,8 +596,17 @@ computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile,
     auto m = map.getSubMap({r});
     LLVM_DEBUG(llvm::dbgs() << "computeSliceParameters: submap: " << m << "\n");
     IRRewriter rewriter(builder);
-    OpFoldResult offset = makeComposedFoldedAffineApply(rewriter, loc, m, lbs);
+    // The offset of the slice is map(lbs) - map(0).
+    SmallVector<Attribute> zeros(lbs.size(), rewriter.getIndexAttr(0));
+    SmallVector<Attribute> mAtZero;
+    auto res = m.constantFold(zeros, mAtZero);
+    assert(succeeded(res));
+    auto atZeroInt = getConstantIntValue(mAtZero[0]);
+    assert(atZeroInt);
+    OpFoldResult offset = makeComposedFoldedAffineApply(
+        rewriter, loc, m.getResult(0) - *atZeroInt, lbs);
     sliceParams.offsets.push_back(offset);
+
     OpFoldResult closedIntSize =
         makeComposedFoldedAffineApply(rewriter, loc, m, subShapeSizes);
     // Resulting size needs to be made half open interval again.
diff --git a/mlir/test/Dialect/Linalg/tile-offset.mlir b/mlir/test/Dialect/Linalg/tile-offset.mlir
new file mode 100644
index 00000000000000..41763dd688c1e9
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/tile-offset.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-opt %s -transform-interpreter -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @tile_offset
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]:
+func.func @tile_offset(%arg0 : tensor<9xf32>) -> tensor<6xf32> {
+  %empty = tensor.empty() : tensor<6xf32>
+  // CHECK: scf.for %[[ITER:[a-zA-Z0-9_]+]] =
+  // CHECK: tensor.extract_slice %[[ARG0]][%[[ITER]]] [6] [1]
+  %generic = linalg.generic
+    {indexing_maps = [affine_map<(d0) -> (d0 + 3)>,
+                      affine_map<(d0) -> (d0)>],
+     iterator_types = ["parallel"]} ins(%arg0: tensor<9xf32>) outs(%empty : tensor<6xf32>) {
+    ^bb0(%in : f32, %out: f32):
+      linalg.yield %in : f32
+    } -> tensor<6xf32>
+  return %generic : tensor<6xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %1, %loop = transform.structured.tile_using_for %0 tile_sizes [3] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+    transform.yield
+  }
+}

>From a01658f80984070ab2c9f6e27c4d120ccc5f8636 Mon Sep 17 00:00:00 2001
From: Matthias Gehre <matthias.gehre at amd.com>
Date: Mon, 13 Jan 2025 09:57:31 +0100
Subject: [PATCH 2/3] Add assert message; make integer extraction more direct

---
 mlir/lib/Dialect/Linalg/Utils/Utils.cpp | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
index 2ebeccb1505048..43e6bbd2c37860 100644
--- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
@@ -596,15 +596,15 @@ computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile,
     auto m = map.getSubMap({r});
     LLVM_DEBUG(llvm::dbgs() << "computeSliceParameters: submap: " << m << "\n");
     IRRewriter rewriter(builder);
-    // The offset of the slice is map(lbs) - map(0).
+    // The offset of the slice is m(lbs) - m(0).
     SmallVector<Attribute> zeros(lbs.size(), rewriter.getIndexAttr(0));
     SmallVector<Attribute> mAtZero;
     auto res = m.constantFold(zeros, mAtZero);
-    assert(succeeded(res));
-    auto atZeroInt = getConstantIntValue(mAtZero[0]);
-    assert(atZeroInt);
+    assert(succeeded(res) && "affine_map must be evaluatable (not symbols)");
+    int64_t mAtZeroInt =
+        cast<IntegerAttr>(mAtZero[0]).getValue().getSExtValue();
     OpFoldResult offset = makeComposedFoldedAffineApply(
-        rewriter, loc, m.getResult(0) - *atZeroInt, lbs);
+        rewriter, loc, m.getResult(0) - mAtZeroInt, lbs);
     sliceParams.offsets.push_back(offset);
 
     OpFoldResult closedIntSize =

>From e0201e9513a452067a41be26a43dbc0677e0e5ff Mon Sep 17 00:00:00 2001
From: Matthias Gehre <matthias.gehre at amd.com>
Date: Mon, 13 Jan 2025 09:59:56 +0100
Subject: [PATCH 3/3] Add  [[maybe_unused]] to bool result

---
 mlir/lib/Dialect/Linalg/Utils/Utils.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
index 43e6bbd2c37860..d148067fe63433 100644
--- a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp
@@ -599,7 +599,7 @@ computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile,
     // The offset of the slice is m(lbs) - m(0).
     SmallVector<Attribute> zeros(lbs.size(), rewriter.getIndexAttr(0));
     SmallVector<Attribute> mAtZero;
-    auto res = m.constantFold(zeros, mAtZero);
+    [[maybe_unused]] auto res = m.constantFold(zeros, mAtZero);
     assert(succeeded(res) && "affine_map must be evaluatable (not symbols)");
     int64_t mAtZeroInt =
         cast<IntegerAttr>(mAtZero[0]).getValue().getSExtValue();