Merge pull request #442 from Xilinx/matthias.fix_tiling_slice_offset

computeSliceParameters: Fix offset when m(0) != 0

Author: mgehre-amd

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

@@ -599,8 +599,17 @@ computeSliceParameters(OpBuilder &builder, Location loc, Value valueToTile,
     // (i.e. the op does not subsample, stepping occurs in the loop).
     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.

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
+  }
+}