diff --git a/mlir/lib/Dialect/Linalg/Utils/Utils.cpp b/mlir/lib/Dialect/Linalg/Utils/Utils.cpp index 38e427af1c484..d148067fe6343 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 m(lbs) - m(0). + SmallVector zeros(lbs.size(), rewriter.getIndexAttr(0)); + SmallVector mAtZero; + [[maybe_unused]] auto res = m.constantFold(zeros, mAtZero); + assert(succeeded(res) && "affine_map must be evaluatable (not symbols)"); + int64_t mAtZeroInt = + cast(mAtZero[0]).getValue().getSExtValue(); + OpFoldResult offset = makeComposedFoldedAffineApply( + rewriter, loc, m.getResult(0) - mAtZeroInt, 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 0000000000000..41763dd688c1e --- /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 + } +}