[Mlir-commits] [mlir] [mlir][tensor] Recognize loop and integer alignment when bubbling reshape slices (PR #236185)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 27 10:01:14 PDT 2026
https://github.com/meinie0826 updated https://github.com/llvm/llvm-project/pull/236185
>From 318f56ad6b8fb097764608ed4a40efcf6f627292 Mon Sep 17 00:00:00 2001
From: meinie0826 <929449726 at qq.com>
Date: Wed, 28 Oct 2026 00:34:25 +0800
Subject: [PATCH 1/2] [mlir][tensor] Recognize power-of-two alignment through
integer expressions and loops
---
.../Dialect/Tensor/Transforms/CMakeLists.txt | 2 +
.../Tensor/Transforms/ReshapePatterns.cpp | 82 ++++-
.../Dialect/Linalg/transform-op-fuse.mlir | 37 +++
.../Tensor/bubble-up-extract-slice-op.mlir | 280 ++++++++++++++++++
4 files changed, 398 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
index 78e64c82544e09..101083a98f8b57 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
@@ -28,7 +28,9 @@ add_mlir_dialect_library(MLIRTensorTransforms
MLIRDialectUtils
MLIRFuncDialect # Dependency of scalarize-single-element-tensor-return
MLIRIR
+ MLIRInferIntDivisibilityOpInterface
MLIRLinalgDialect
+ MLIRLoopLikeInterface
MLIRMemRefDialect
MLIRPass
MLIRRuntimeVerifiableOpInterface
diff --git a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
index c9bd4181da4f84..b12792c363a22a 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
@@ -7,14 +7,20 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Value.h"
+#include "mlir/Interfaces/InferIntDivisibilityOpInterface.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/LogicalResult.h"
+#include "llvm/Support/MathExtras.h"
+
+#include <numeric>
using namespace mlir;
using namespace mlir::tensor;
@@ -618,14 +624,84 @@ LogicalResult mlir::tensor::getCollapsedExtractSliceInfo(
return success();
}
-// Checks if the `ofr` is a multiple of the `factor`.
-// Handles both static integer and dynamic values
-// where the value is the result of an affine.apply.
+/// Infer a power-of-two divisor from integer expressions and induction variable
+/// bounds. The query is local to a match: no facts are retained across
+/// rewrites. The budget bounds work on shared expression DAGs and nested loop
+/// bounds.
+static uint64_t inferKnownPowerOfTwoDivisor(OpFoldResult ofr,
+ unsigned &budget) {
+ if (!budget)
+ return 1;
+ --budget;
+
+ if (auto constant = getConstantIntValue(ofr)) {
+ uint64_t bits = static_cast<uint64_t>(*constant);
+ // Keep only the power-of-two part, valid for both signed and unsigned
+ // interpretations. Zero remains divisible by every factor.
+ return bits & (uint64_t{0} - bits);
+ }
+
+ Value value = cast<Value>(ofr);
+ if (auto arg = dyn_cast<BlockArgument>(value)) {
+ auto loop = dyn_cast<LoopLikeOpInterface>(arg.getOwner()->getParentOp());
+ if (!loop)
+ return 1;
+ auto ivs = loop.getLoopInductionVars();
+ auto lbs = loop.getLoopLowerBounds();
+ auto steps = loop.getLoopSteps();
+ if (!ivs || !lbs || !steps)
+ return 1;
+ auto it = llvm::find(*ivs, value);
+ if (it == ivs->end())
+ return 1;
+ unsigned index = std::distance(ivs->begin(), it);
+ if (index >= lbs->size() || index >= steps->size())
+ return 1;
+ // An induction variable is lb + iteration * step.
+ uint64_t lbDivisor = inferKnownPowerOfTwoDivisor((*lbs)[index], budget);
+ uint64_t stepDivisor = inferKnownPowerOfTwoDivisor((*steps)[index], budget);
+ return std::gcd(lbDivisor, stepDivisor);
+ }
+
+ Operation *op = value.getDefiningOp();
+ // For these operations, power-of-two divisibility is preserved even when
+ // arithmetic wraps. Other operations are left to the affine proof below.
+ if (!isa<arith::AddIOp, arith::SubIOp, arith::MulIOp, arith::SelectOp>(op))
+ return 1;
+ auto inferable = cast<InferIntDivisibilityOpInterface>(op);
+ SmallVector<IntegerDivisibility> operands;
+ for (Value operand : op->getOperands()) {
+ uint64_t divisor = inferKnownPowerOfTwoDivisor(operand, budget);
+ operands.emplace_back(ConstantIntDivisibility(divisor, divisor));
+ }
+ uint64_t divisor = 1;
+ inferable.inferResultDivisibility(
+ operands, [&](Value result, const ConstantIntDivisibility &inferred) {
+ if (result != value)
+ return;
+ // Saturate overflowing divisor products rather than treating them
+ // as a proof that the result is exactly zero.
+ divisor = inferred.udiv() ? inferred.udiv() : uint64_t{1} << 63;
+ });
+ return divisor;
+}
+
+// Check static and affine divisibility, or power-of-two alignment inferred
+// from integer expressions and loop bounds. Restricting the latter proof to
+// powers of two makes it independent of the index bit width and wraparound.
static bool isMultipleOf(OpFoldResult ofr, int64_t factor) {
+ if (factor <= 0)
+ return false;
std::optional<int64_t> staticValue = getConstantIntValue(ofr);
if (staticValue.has_value())
return staticValue.value() % factor == 0;
+ if (llvm::isPowerOf2_64(factor)) {
+ unsigned budget = 32;
+ if (inferKnownPowerOfTwoDivisor(ofr, budget) % factor == 0)
+ return true;
+ }
+
Value value = dyn_cast<Value>(ofr);
if (!value)
return false;
diff --git a/mlir/test/Dialect/Linalg/transform-op-fuse.mlir b/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
index 538e5765cdc230..0a5aa41a89caa8 100644
--- a/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
@@ -899,3 +899,40 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// CHECK-LABEL: func.func @fuse_matmul_through_collapse(
+// CHECK-NOT: linalg.matmul
+// CHECK: scf.for
+// CHECK: %[[INIT:.*]] = linalg.fill
+// CHECK-SAME: -> tensor<1x8xf32>
+// CHECK: %[[MM:.*]] = linalg.matmul
+// CHECK-SAME: outs(%[[INIT]] : tensor<1x8xf32>) -> tensor<1x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[MM]]
+// CHECK: linalg.generic
+// CHECK-SAME: ins(%[[FLAT]] : tensor<8xf32>)
+// CHECK: tensor.insert_slice
+func.func @fuse_matmul_through_collapse(%a: tensor<8x4xf32>, %b: tensor<4x8xf32>) -> tensor<64xf32> {
+ %zero = arith.constant 0.0 : f32
+ %empty = tensor.empty() : tensor<8x8xf32>
+ %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<8x8xf32>) -> tensor<8x8xf32>
+ %mm = linalg.matmul ins(%a, %b : tensor<8x4xf32>, tensor<4x8xf32>) outs(%init : tensor<8x8xf32>) -> tensor<8x8xf32>
+ %flat = tensor.collapse_shape %mm [[0, 1]] : tensor<8x8xf32> into tensor<64xf32>
+ %out = tensor.empty() : tensor<64xf32>
+ %r = linalg.generic {candidate_root, indexing_maps = [affine_map<(i)->(i)>, affine_map<(i)->(i)>], iterator_types = ["parallel"]}
+ ins(%flat : tensor<64xf32>) outs(%out : tensor<64xf32>) {
+ ^bb0(%v: f32, %old: f32):
+ %v1 = arith.maximumf %v, %zero : f32
+ linalg.yield %v1 : f32
+ } -> tensor<64xf32>
+ return %r : tensor<64xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %f = transform.structured.match attributes{candidate_root} in %root : (!transform.any_op) -> !transform.any_op
+ %t, %l = transform.structured.fuse %f tile_sizes [8] apply_cleanup : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
diff --git a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
index 31232ee7151b2b..3490782e188d6a 100644
--- a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
+++ b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
@@ -492,3 +492,283 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// Power-of-two alignment through integer expressions and loop bounds.
+
+// CHECK-LABEL: func.func @bubble_collapse_muli_offset(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_muli_offset(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %offset = arith.muli %row, %c8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_addi_offset(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_addi_offset(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %base = arith.muli %row, %c8 : index
+ %offset = arith.addi %base, %c8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_subi_offset(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_subi_offset(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %base = arith.muli %row, %c8 : index
+ %offset = arith.subi %base, %c8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_select_offset(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_select_offset(%src: tensor<4x8xf32>, %row: index, %cond: i1) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %a = arith.muli %row, %c8 : index
+ %b = arith.addi %a, %c8 : index
+ %offset = arith.select %cond, %a, %b : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_collapse_misaligned_add(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_misaligned_add(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %c1 = arith.constant 1 : index
+ %base = arith.muli %row, %c8 : index
+ %offset = arith.addi %base, %c1 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_collapse_misaligned_select(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_misaligned_select(%src: tensor<4x8xf32>, %row: index, %cond: i1) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %c1 = arith.constant 1 : index
+ %a = arith.muli %row, %c8 : index
+ %b = arith.addi %a, %c1 : index
+ %offset = arith.select %cond, %a, %b : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_collapse_non_power_of_two_muli(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_non_power_of_two_muli(%src: tensor<4x6xf32>, %row: index) -> tensor<6xf32> {
+ %c6 = arith.constant 6 : index
+ %offset = arith.muli %row, %c6 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %slice = tensor.extract_slice %flat[%offset] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ return %slice : tensor<6xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_collapse_unsupported_division(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_unsupported_division(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %c2 = arith.constant 2 : index
+ %base = arith.muli %row, %c8 : index
+ %offset = arith.divui %base, %c2 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_collapse_query_budget(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_query_budget(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %v0 = arith.muli %row, %c8 : index
+ %v1 = arith.addi %v0, %v0 : index
+ %v2 = arith.addi %v1, %v1 : index
+ %v3 = arith.addi %v2, %v2 : index
+ %v4 = arith.addi %v3, %v3 : index
+ %v5 = arith.addi %v4, %v4 : index
+ %v6 = arith.addi %v5, %v5 : index
+ %offset = arith.addi %v6, %v6 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_multiple_groups(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8x2x4xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8, 1, 4] [1, 1, 1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_multiple_groups(%src: tensor<4x8x2x4xf32>, %row: index, %col: index) -> tensor<8x4xf32> {
+ %c8 = arith.constant 8 : index
+ %c4 = arith.constant 4 : index
+ %a = arith.muli %row, %c8 : index
+ %b = arith.muli %col, %c4 : index
+ %flat = tensor.collapse_shape %src [[0, 1], [2, 3]] : tensor<4x8x2x4xf32> into tensor<32x8xf32>
+ %slice = tensor.extract_slice %flat[%a, %b] [8, 4] [1, 1] : tensor<32x8xf32> to tensor<8x4xf32>
+ return %slice : tensor<8x4xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_loop_iv(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: scf.for
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_loop_iv(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
+ %lb = arith.constant 0 : index
+ %ub = arith.constant 25 : index
+ %step = arith.constant 8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %r = scf.for %iv = %lb to %ub step %step iter_args(%out = %dst) -> tensor<32xf32> {
+ %slice = tensor.extract_slice %flat[%iv] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ %insert = tensor.insert_slice %slice into %out[%iv] [8] [1] : tensor<8xf32> into tensor<32xf32>
+ scf.yield %insert : tensor<32xf32>
+ }
+ return %r : tensor<32xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_aligned_loop_lb(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: scf.for
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_aligned_loop_lb(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
+ %lb = arith.constant 8 : index
+ %ub = arith.constant 25 : index
+ %step = arith.constant 8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %r = scf.for %iv = %lb to %ub step %step iter_args(%out = %dst) -> tensor<32xf32> {
+ %slice = tensor.extract_slice %flat[%iv] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ %insert = tensor.insert_slice %slice into %out[%iv] [8] [1] : tensor<8xf32> into tensor<32xf32>
+ scf.yield %insert : tensor<32xf32>
+ }
+ return %r : tensor<32xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_collapse_misaligned_loop_lb(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: scf.for
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_misaligned_loop_lb(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
+ %lb = arith.constant 1 : index
+ %ub = arith.constant 25 : index
+ %step = arith.constant 8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %r = scf.for %iv = %lb to %ub step %step iter_args(%out = %dst) -> tensor<32xf32> {
+ %slice = tensor.extract_slice %flat[%iv] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ %insert = tensor.insert_slice %slice into %out[%iv] [8] [1] : tensor<8xf32> into tensor<32xf32>
+ scf.yield %insert : tensor<32xf32>
+ }
+ return %r : tensor<32xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_collapse_misaligned_loop_step(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: scf.for
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_misaligned_loop_step(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
+ %lb = arith.constant 0 : index
+ %ub = arith.constant 25 : index
+ %step = arith.constant 4 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %r = scf.for %iv = %lb to %ub step %step iter_args(%out = %dst) -> tensor<32xf32> {
+ %slice = tensor.extract_slice %flat[%iv] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ %insert = tensor.insert_slice %slice into %out[%iv] [8] [1] : tensor<8xf32> into tensor<32xf32>
+ scf.yield %insert : tensor<32xf32>
+ }
+ return %r : tensor<32xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_collapse_iter_arg(
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape
+// CHECK: scf.for
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_iter_arg(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %r:2 = scf.for %iv = %c0 to %c4 step %c1 iter_args(%offset = %c0, %out = %dst) -> (index, tensor<32xf32>) {
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ %insert = tensor.insert_slice %slice into %out[%offset] [8] [1] : tensor<8xf32> into tensor<32xf32>
+ %next = arith.addi %offset, %c8 : index
+ scf.yield %next, %insert : index, tensor<32xf32>
+ }
+ return %r#1 : tensor<32xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_negative_multiplier(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_negative_multiplier(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %cneg8 = arith.constant -8 : index
+ %offset = arith.muli %row, %cneg8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_divisor_product_overflow(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_divisor_product_overflow(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c = arith.constant -9223372036854775808 : index
+ %c2 = arith.constant 2 : index
+ %a = arith.muli %row, %c : index
+ %offset = arith.muli %a, %c2 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.tensor.bubble_up_extract_slice
+ } : !transform.any_op
+ transform.yield
+ }
+}
>From d1c990e8b173e7318d1c1ba9763786a26292cc61 Mon Sep 17 00:00:00 2001
From: meinie0826 <929449726 at qq.com>
Date: Wed, 28 Oct 2026 01:00:53 +0800
Subject: [PATCH 2/2] [mlir][tensor] Use scoped dataflow analysis for reshape
fusion cleanup
---
.../SCF/Transforms/TileUsingInterface.h | 7 +
.../Dialect/Tensor/Transforms/Transforms.h | 28 +-
.../TransformOps/LinalgTransformOps.cpp | 21 +-
.../SCF/Transforms/TileUsingInterface.cpp | 17 +-
.../Dialect/Tensor/Transforms/CMakeLists.txt | 1 +
.../Tensor/Transforms/ReshapePatterns.cpp | 233 +++++++----
.../bubble-up-extract-slice-alignment.mlir | 367 ++++++++++++++++++
.../Tensor/bubble-up-extract-slice-op.mlir | 280 -------------
.../Dialect/Tensor/TestTensorTransforms.cpp | 34 ++
9 files changed, 610 insertions(+), 378 deletions(-)
create mode 100644 mlir/test/Dialect/Tensor/bubble-up-extract-slice-alignment.mlir
diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
index f457da8b3d71df..10a616428c2619 100644
--- a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
+++ b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
@@ -316,6 +316,13 @@ struct SCFTileAndFuseOptions {
/// before fusion. This will track deleted and newly inserted
/// `tensor.extract_slice` ops and update the worklist.
std::optional<FrozenRewritePatternSet> cleanupPatterns = std::nullopt;
+
+ /// Optional cleanup callback invoked after tiling and after each fusion step,
+ /// instead of `cleanupPatterns`. This permits analysis-backed cleanup to
+ /// analyze the generated IR at each phase. All rewrites must notify
+ /// `listener` so the extract_slice worklist remains up to date.
+ std::function<LogicalResult(ArrayRef<Operation *>, RewriterBase::Listener *)>
+ cleanupFn;
};
/// Result of fusing the producer of the source of a `tensor.extract_slice`.
diff --git a/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
index 341b3b1dc3472e..552ea7074d24cc 100644
--- a/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
@@ -70,6 +70,16 @@ void populateMergeConsecutiveInsertExtractSlicePatterns(
/// tiling interface.
void populateBubbleUpExtractSliceOpPatterns(RewritePatternSet &patterns);
+/// Bubble up the extract_slice operations in `operations` through their
+/// collapse_shape producers. Analyze integer alignment in `analysisScope` once,
+/// complete matching without modifying existing definitions, then replace the
+/// slices. Return the surviving operations and the inserted slice/reshape ops
+/// for subsequent cleanup. The analysis is not retained across invocations.
+FailureOr<SmallVector<Operation *>>
+bubbleUpExtractSlicesThroughCollapseShape(RewriterBase &rewriter,
+ Operation *analysisScope,
+ ArrayRef<Operation *> operations);
+
/// Populates `patterns` with patterns that drop redundant tensor.insert_slice
/// rank expansions.
void populateDropRedundantInsertSliceRankExpansionPatterns(
@@ -161,15 +171,15 @@ getCollapsedExtractSliceInfo(OpBuilder &b, tensor::ExtractSliceOp sliceOp,
/// the shape of `expandedValue`.
///
/// This fails when the specified expansion cannot be represented by a valid
-/// ExtractSliceOp.
-LogicalResult
-getExpandedExtractSliceInfo(OpBuilder &b, tensor::ExtractSliceOp sliceOp,
- ArrayRef<ReassociationIndices> reassociation,
- Value expandedValue,
- SmallVectorImpl<OpFoldResult> &expandedOffsets,
- SmallVectorImpl<OpFoldResult> &expandedSizes,
- SmallVectorImpl<OpFoldResult> &expandedStrides);
-
+/// ExtractSliceOp. `isAligned` may provide additional divisibility proofs for
+/// offsets and must remain valid during this call.
+LogicalResult getExpandedExtractSliceInfo(
+ OpBuilder &b, tensor::ExtractSliceOp sliceOp,
+ ArrayRef<ReassociationIndices> reassociation, Value expandedValue,
+ SmallVectorImpl<OpFoldResult> &expandedOffsets,
+ SmallVectorImpl<OpFoldResult> &expandedSizes,
+ SmallVectorImpl<OpFoldResult> &expandedStrides,
+ llvm::function_ref<bool(Value, int64_t)> isAligned = nullptr);
} // namespace tensor
} // namespace mlir
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 816d7499845789..b82f37ac59a9f4 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -753,7 +753,26 @@ transform::FuseOp::apply(transform::TransformRewriter &rewriter,
tensor::ExtractSliceOp::getCanonicalizationPatterns(patterns, context);
tensor::populateMergeConsecutiveInsertExtractSlicePatterns(patterns);
tensor::populateBubbleUpExtractSliceOpPatterns(patterns);
- tileAndFuseOptions.cleanupPatterns = std::move(patterns);
+ tileAndFuseOptions.cleanupFn =
+ [context, patterns = FrozenRewritePatternSet(std::move(patterns))](
+ ArrayRef<Operation *> ops, RewriterBase::Listener *listener) {
+ // Tiling and fusion can create or replace loops. Analyze the
+ // enclosing isolated scope for each batch rather than keeping
+ // pre-tiling facts.
+ Operation *scope = ops.front();
+ while (!scope->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
+ scope->getParentOp())
+ scope = scope->getParentOp();
+ IRRewriter cleanupRewriter(context, listener);
+ auto cleaned = tensor::bubbleUpExtractSlicesThroughCollapseShape(
+ cleanupRewriter, scope, ops);
+ if (failed(cleaned))
+ return failure();
+ return applyOpPatternsGreedily(
+ *cleaned, patterns,
+ GreedyRewriteConfig().setListener(listener).setStrictness(
+ GreedyRewriteStrictness::ExistingAndNewOps));
+ };
}
size_t numLoops;
diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index a8efe6cdb4415e..57b7428098de76 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -1619,8 +1619,7 @@ namespace {
/// fusion algorithm to apply cleanup patterns in between fusion steps.
class SliceTrackingListener : public RewriterBase::Listener {
public:
- explicit SliceTrackingListener(
- std::optional<FrozenRewritePatternSet> patterns);
+ explicit SliceTrackingListener(const scf::SCFTileAndFuseOptions &options);
SliceTrackingListener() = default;
/// Adds the given list of operations to the worklist, and if present,
@@ -1650,11 +1649,14 @@ class SliceTrackingListener : public RewriterBase::Listener {
/// Optional pattern set to apply when adding new operations to the
/// worklist.
std::optional<FrozenRewritePatternSet> patterns = std::nullopt;
+ std::function<LogicalResult(ArrayRef<Operation *>, RewriterBase::Listener *)>
+ cleanupFn;
};
SliceTrackingListener::SliceTrackingListener(
- std::optional<FrozenRewritePatternSet> p) {
- patterns = std::move(p);
+ const scf::SCFTileAndFuseOptions &options) {
+ patterns = options.cleanupPatterns;
+ cleanupFn = options.cleanupFn;
}
LogicalResult
@@ -1664,6 +1666,10 @@ SliceTrackingListener::insertAndApplyPatterns(ArrayRef<Operation *> ops) {
worklist.push_back(slice);
}
+ if (ops.empty())
+ return success();
+ if (cleanupFn)
+ return cleanupFn(ops, this);
if (!patterns)
return success();
@@ -1806,8 +1812,7 @@ mlir::scf::tileConsumerAndFuseProducersUsingSCF(
SCFTileAndFuseOptions::ControlFnResult controlFnResult;
};
- SliceTrackingListener sliceTracker =
- SliceTrackingListener(options.cleanupPatterns);
+ SliceTrackingListener sliceTracker = SliceTrackingListener(options);
if (failed(
sliceTracker.insertAndApplyPatterns(tilingResult->generatedSlices))) {
diff --git a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
index 101083a98f8b57..9198e307017a14 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
@@ -20,6 +20,7 @@ add_mlir_dialect_library(MLIRTensorTransforms
MLIRTensorTransformsIncGen
LINK_LIBS PUBLIC
+ MLIRAnalysis
MLIRAffineDialect
MLIRAffineTransforms
MLIRAffineUtils
diff --git a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
index b12792c363a22a..6f2a3705392723 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
@@ -6,6 +6,9 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"
+#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
+#include "mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
@@ -17,11 +20,10 @@
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/MathExtras.h"
-#include <numeric>
-
using namespace mlir;
using namespace mlir::tensor;
@@ -624,87 +626,20 @@ LogicalResult mlir::tensor::getCollapsedExtractSliceInfo(
return success();
}
-/// Infer a power-of-two divisor from integer expressions and induction variable
-/// bounds. The query is local to a match: no facts are retained across
-/// rewrites. The budget bounds work on shared expression DAGs and nested loop
-/// bounds.
-static uint64_t inferKnownPowerOfTwoDivisor(OpFoldResult ofr,
- unsigned &budget) {
- if (!budget)
- return 1;
- --budget;
-
- if (auto constant = getConstantIntValue(ofr)) {
- uint64_t bits = static_cast<uint64_t>(*constant);
- // Keep only the power-of-two part, valid for both signed and unsigned
- // interpretations. Zero remains divisible by every factor.
- return bits & (uint64_t{0} - bits);
- }
-
- Value value = cast<Value>(ofr);
- if (auto arg = dyn_cast<BlockArgument>(value)) {
- auto loop = dyn_cast<LoopLikeOpInterface>(arg.getOwner()->getParentOp());
- if (!loop)
- return 1;
- auto ivs = loop.getLoopInductionVars();
- auto lbs = loop.getLoopLowerBounds();
- auto steps = loop.getLoopSteps();
- if (!ivs || !lbs || !steps)
- return 1;
- auto it = llvm::find(*ivs, value);
- if (it == ivs->end())
- return 1;
- unsigned index = std::distance(ivs->begin(), it);
- if (index >= lbs->size() || index >= steps->size())
- return 1;
- // An induction variable is lb + iteration * step.
- uint64_t lbDivisor = inferKnownPowerOfTwoDivisor((*lbs)[index], budget);
- uint64_t stepDivisor = inferKnownPowerOfTwoDivisor((*steps)[index], budget);
- return std::gcd(lbDivisor, stepDivisor);
- }
-
- Operation *op = value.getDefiningOp();
- // For these operations, power-of-two divisibility is preserved even when
- // arithmetic wraps. Other operations are left to the affine proof below.
- if (!isa<arith::AddIOp, arith::SubIOp, arith::MulIOp, arith::SelectOp>(op))
- return 1;
- auto inferable = cast<InferIntDivisibilityOpInterface>(op);
- SmallVector<IntegerDivisibility> operands;
- for (Value operand : op->getOperands()) {
- uint64_t divisor = inferKnownPowerOfTwoDivisor(operand, budget);
- operands.emplace_back(ConstantIntDivisibility(divisor, divisor));
- }
- uint64_t divisor = 1;
- inferable.inferResultDivisibility(
- operands, [&](Value result, const ConstantIntDivisibility &inferred) {
- if (result != value)
- return;
- // Saturate overflowing divisor products rather than treating them
- // as a proof that the result is exactly zero.
- divisor = inferred.udiv() ? inferred.udiv() : uint64_t{1} << 63;
- });
- return divisor;
-}
-
-// Check static and affine divisibility, or power-of-two alignment inferred
-// from integer expressions and loop bounds. Restricting the latter proof to
-// powers of two makes it independent of the index bit width and wraparound.
-static bool isMultipleOf(OpFoldResult ofr, int64_t factor) {
+// Check static and affine divisibility, or consult the caller's alignment
+// proof. The query must remain valid until all slice information is computed.
+static bool isMultipleOf(OpFoldResult ofr, int64_t factor,
+ llvm::function_ref<bool(Value, int64_t)> isAligned) {
if (factor <= 0)
return false;
- std::optional<int64_t> staticValue = getConstantIntValue(ofr);
- if (staticValue.has_value())
- return staticValue.value() % factor == 0;
-
- if (llvm::isPowerOf2_64(factor)) {
- unsigned budget = 32;
- if (inferKnownPowerOfTwoDivisor(ofr, budget) % factor == 0)
- return true;
- }
+ if (auto constant = getConstantIntValue(ofr))
+ return *constant % factor == 0;
Value value = dyn_cast<Value>(ofr);
if (!value)
return false;
+ if (isAligned && isAligned(value, factor))
+ return true;
auto applyOp = value.getDefiningOp<affine::AffineApplyOp>();
if (!applyOp)
return false;
@@ -726,7 +661,8 @@ static bool isMultipleOf(OpFoldResult ofr, int64_t factor) {
static LogicalResult computeExpandedSliceInfoForReassocGroup(
OpBuilder &b, OpFoldResult collapsedSize, OpFoldResult collapsedOffset,
const ReassociationIndices &reassocIndices, ArrayRef<int64_t> expandedShape,
- SmallVectorImpl<OpFoldResult> &groupSizes) {
+ SmallVectorImpl<OpFoldResult> &groupSizes,
+ llvm::function_ref<bool(Value, int64_t)> isAligned) {
assert(groupSizes.empty() && "Group sizes must be empty");
// The first case is when there's only one non-unit dimension in the
// reassociation group.
@@ -808,7 +744,7 @@ static LogicalResult computeExpandedSliceInfoForReassocGroup(
// Check dynamic/static offset divisibility.
currentOffsetDivisor *= expandedShapeSize;
- if (!isMultipleOf(collapsedOffset, currentOffsetDivisor))
+ if (!isMultipleOf(collapsedOffset, currentOffsetDivisor, isAligned))
return failure();
// Trailing dims get full shape and zero offset.
@@ -836,7 +772,7 @@ static LogicalResult computeExpandedSliceInfoForReassocGroup(
// to check the validity of the range.
if ((expandedShapeSize % currentCollapsedsize) != 0)
return failure();
- if (!isMultipleOf(collapsedOffset, staticSize.value()))
+ if (!isMultipleOf(collapsedOffset, staticSize.value(), isAligned))
return failure();
}
// Slicing dimension gets the remaining collapsed size.
@@ -863,7 +799,8 @@ LogicalResult mlir::tensor::getExpandedExtractSliceInfo(
ArrayRef<ReassociationIndices> reassociation, Value expandedValue,
SmallVectorImpl<OpFoldResult> &expandedOffsets,
SmallVectorImpl<OpFoldResult> &expandedSizes,
- SmallVectorImpl<OpFoldResult> &expandedStrides) {
+ SmallVectorImpl<OpFoldResult> &expandedStrides,
+ llvm::function_ref<bool(Value, int64_t)> isAligned) {
if (!sliceOp.hasUnitStride()) {
return failure();
}
@@ -893,7 +830,7 @@ LogicalResult mlir::tensor::getExpandedExtractSliceInfo(
SmallVector<OpFoldResult> groupSizes;
LogicalResult result = computeExpandedSliceInfoForReassocGroup(
b, collapsedSize, collapsedOffset, reassocIndices, expandedShape,
- groupSizes);
+ groupSizes, isAligned);
if (failed(result))
return failure();
groupResults.emplace_back(groupSizes);
@@ -920,6 +857,138 @@ LogicalResult mlir::tensor::getExpandedExtractSliceInfo(
return success();
}
+namespace {
+/// Restrict operation transfer functions to facts that survive modular integer
+/// arithmetic. Reuse the existing analysis for the fixed point, control flow,
+/// and induction variables; do not propagate general divisors through division.
+class SliceAlignmentAnalysis : public dataflow::IntegerDivisibilityAnalysis {
+public:
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SliceAlignmentAnalysis)
+ using IntegerDivisibilityAnalysis::IntegerDivisibilityAnalysis;
+
+ LogicalResult visitOperation(
+ Operation *op,
+ ArrayRef<const dataflow::IntegerDivisibilityLattice *> operands,
+ ArrayRef<dataflow::IntegerDivisibilityLattice *> results) override {
+ if (auto constant = dyn_cast<arith::ConstantOp>(op)) {
+ auto attr = dyn_cast<IntegerAttr>(constant.getValue());
+ if (!attr) {
+ setAllToEntryStates(results);
+ return success();
+ }
+ const APInt &value = attr.getValue();
+ // Avoid signed absolute values and support constants wider than uint64_t.
+ uint64_t divisor =
+ value.isZero() ? 0
+ : uint64_t{1} << std::min(value.countr_zero(), 63u);
+ propagateIfChanged(
+ results.front(),
+ results.front()->join(ConstantIntDivisibility(divisor, divisor)));
+ return success();
+ }
+ if (!isa<arith::AddIOp, arith::SubIOp, arith::MulIOp, arith::SelectOp>(
+ op)) {
+ setAllToEntryStates(results);
+ return success();
+ }
+ if (llvm::any_of(operands, [](const auto *lattice) {
+ return lattice->getValue().isUninitialized();
+ }))
+ return success();
+ return IntegerDivisibilityAnalysis::visitOperation(op, operands, results);
+ }
+
+ void visitNonControlFlowArguments(
+ Operation *op, const RegionSuccessor &successor,
+ ValueRange successorInputs,
+ ArrayRef<dataflow::IntegerDivisibilityLattice *> argLattices) override {
+ // The existing loop-bound query represents constants in int64_t.
+ if (auto loop = dyn_cast<LoopLikeOpInterface>(op)) {
+ auto ivs = loop.getLoopInductionVars();
+ if (ivs && llvm::any_of(*ivs, [](Value iv) {
+ auto type = dyn_cast<IntegerType>(iv.getType());
+ return type && type.getWidth() > 64;
+ })) {
+ return SparseForwardDataFlowAnalysis::visitNonControlFlowArguments(
+ op, successor, successorInputs, argLattices);
+ }
+ }
+ IntegerDivisibilityAnalysis::visitNonControlFlowArguments(
+ op, successor, successorInputs, argLattices);
+ }
+};
+} // namespace
+
+FailureOr<SmallVector<Operation *>>
+mlir::tensor::bubbleUpExtractSlicesThroughCollapseShape(
+ RewriterBase &rewriter, Operation *analysisScope,
+ ArrayRef<Operation *> operations) {
+ if (llvm::none_of(operations, [](Operation *op) {
+ auto slice = dyn_cast<ExtractSliceOp>(op);
+ return slice && slice.getSource().getDefiningOp<CollapseShapeOp>();
+ }))
+ return llvm::to_vector(operations);
+
+ DataFlowSolver solver;
+ solver.load<dataflow::DeadCodeAnalysis>();
+ solver.load<dataflow::SparseConstantPropagation>();
+ solver.load<SliceAlignmentAnalysis>();
+ if (failed(solver.initializeAndRun(analysisScope)))
+ return failure();
+
+ auto isAligned = [&](Value value, int64_t factor) {
+ if (!llvm::isPowerOf2_64(factor))
+ return false;
+ const auto *lattice =
+ solver.lookupState<dataflow::IntegerDivisibilityLattice>(value);
+ if (!lattice || lattice->getValue().isUninitialized())
+ return false;
+ return lattice->getValue().getValue().udiv() % factor == 0;
+ };
+
+ // Complete every analysis query before replacing any operation. Building
+ // slices only inserts operations: existing integer definitions, loop bounds,
+ // and control flow stay unchanged throughout matching.
+ SmallVector<std::pair<Operation *, Operation *>> replacements;
+ SmallVector<Operation *> cleanedOperations;
+ OpBuilder::InsertionGuard guard(rewriter);
+ llvm::SmallPtrSet<Operation *, 16> seen;
+ for (Operation *op : operations) {
+ if (!seen.insert(op).second)
+ continue;
+ auto slice = dyn_cast<ExtractSliceOp>(op);
+ auto collapse = slice ? slice.getSource().getDefiningOp<CollapseShapeOp>()
+ : CollapseShapeOp();
+ if (!collapse) {
+ cleanedOperations.push_back(op);
+ continue;
+ }
+ rewriter.setInsertionPoint(slice);
+ SmallVector<OpFoldResult> offsets, sizes, strides;
+ if (failed(getExpandedExtractSliceInfo(
+ rewriter, slice, collapse.getReassociationIndices(),
+ collapse.getSrc(), offsets, sizes, strides, isAligned))) {
+ cleanedOperations.push_back(op);
+ continue;
+ }
+ auto expandedSlice = ExtractSliceOp::create(
+ rewriter, slice.getLoc(), collapse.getSrc(), offsets, sizes, strides);
+ auto replacement = CollapseShapeOp::create(
+ rewriter, slice.getLoc(), slice.getType(), expandedSlice,
+ collapse.getReassociationIndices());
+ replacements.emplace_back(slice, replacement);
+ // Include the expanded slice so the greedy cleanup can visit its producer.
+ cleanedOperations.push_back(expandedSlice);
+ cleanedOperations.push_back(replacement);
+ }
+
+ // No lattice is queried after this point. The next cleanup phase constructs
+ // a solver for the IR containing any newly fused or replaced loops.
+ for (auto [slice, replacement] : replacements)
+ rewriter.replaceOp(slice, replacement);
+ return cleanedOperations;
+}
+
void mlir::tensor::populateReassociativeReshapeFoldingPatterns(
RewritePatternSet &patterns) {
patterns.add<FoldExpandOfRankReducingExtract, FoldExtractSliceOfExpandShape,
diff --git a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-alignment.mlir b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-alignment.mlir
new file mode 100644
index 00000000000000..1764ea3b3592d1
--- /dev/null
+++ b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-alignment.mlir
@@ -0,0 +1,367 @@
+// RUN: mlir-opt %s -test-tensor-transform-patterns="test-bubble-up-extract-slice-alignment" | FileCheck %s
+
+// Power-of-two alignment through integer expressions and loop bounds.
+
+// CHECK-LABEL: func.func @bubble_collapse_muli_offset(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_muli_offset(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %offset = arith.muli %row, %c8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_addi_offset(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_addi_offset(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %base = arith.muli %row, %c8 : index
+ %offset = arith.addi %base, %c8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_subi_offset(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_subi_offset(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %base = arith.muli %row, %c8 : index
+ %offset = arith.subi %base, %c8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_select_offset(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_select_offset(%src: tensor<4x8xf32>, %row: index, %cond: i1) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %a = arith.muli %row, %c8 : index
+ %b = arith.addi %a, %c8 : index
+ %offset = arith.select %cond, %a, %b : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_collapse_misaligned_add(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_misaligned_add(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %c1 = arith.constant 1 : index
+ %base = arith.muli %row, %c8 : index
+ %offset = arith.addi %base, %c1 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_collapse_misaligned_select(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_misaligned_select(%src: tensor<4x8xf32>, %row: index, %cond: i1) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %c1 = arith.constant 1 : index
+ %a = arith.muli %row, %c8 : index
+ %b = arith.addi %a, %c1 : index
+ %offset = arith.select %cond, %a, %b : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_collapse_non_power_of_two_muli(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_non_power_of_two_muli(%src: tensor<4x6xf32>, %row: index) -> tensor<6xf32> {
+ %c6 = arith.constant 6 : index
+ %offset = arith.muli %row, %c6 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %slice = tensor.extract_slice %flat[%offset] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ return %slice : tensor<6xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_collapse_unsupported_division(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_unsupported_division(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %c2 = arith.constant 2 : index
+ %base = arith.muli %row, %c8 : index
+ %offset = arith.divui %base, %c2 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_shared_expression_dag(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_shared_expression_dag(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c8 = arith.constant 8 : index
+ %v0 = arith.muli %row, %c8 : index
+ %v1 = arith.addi %v0, %v0 : index
+ %v2 = arith.addi %v1, %v1 : index
+ %v3 = arith.addi %v2, %v2 : index
+ %v4 = arith.addi %v3, %v3 : index
+ %v5 = arith.addi %v4, %v4 : index
+ %v6 = arith.addi %v5, %v5 : index
+ %offset = arith.addi %v6, %v6 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_multiple_groups(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8x2x4xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8, 1, 4] [1, 1, 1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_multiple_groups(%src: tensor<4x8x2x4xf32>, %row: index, %col: index) -> tensor<8x4xf32> {
+ %c8 = arith.constant 8 : index
+ %c4 = arith.constant 4 : index
+ %a = arith.muli %row, %c8 : index
+ %b = arith.muli %col, %c4 : index
+ %flat = tensor.collapse_shape %src [[0, 1], [2, 3]] : tensor<4x8x2x4xf32> into tensor<32x8xf32>
+ %slice = tensor.extract_slice %flat[%a, %b] [8, 4] [1, 1] : tensor<32x8xf32> to tensor<8x4xf32>
+ return %slice : tensor<8x4xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_loop_iv(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: scf.for
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_loop_iv(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
+ %lb = arith.constant 0 : index
+ %ub = arith.constant 25 : index
+ %step = arith.constant 8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %r = scf.for %iv = %lb to %ub step %step iter_args(%out = %dst) -> tensor<32xf32> {
+ %slice = tensor.extract_slice %flat[%iv] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ %insert = tensor.insert_slice %slice into %out[%iv] [8] [1] : tensor<8xf32> into tensor<32xf32>
+ scf.yield %insert : tensor<32xf32>
+ }
+ return %r : tensor<32xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_aligned_loop_lb(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: scf.for
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_aligned_loop_lb(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
+ %lb = arith.constant 8 : index
+ %ub = arith.constant 25 : index
+ %step = arith.constant 8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %r = scf.for %iv = %lb to %ub step %step iter_args(%out = %dst) -> tensor<32xf32> {
+ %slice = tensor.extract_slice %flat[%iv] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ %insert = tensor.insert_slice %slice into %out[%iv] [8] [1] : tensor<8xf32> into tensor<32xf32>
+ scf.yield %insert : tensor<32xf32>
+ }
+ return %r : tensor<32xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_collapse_misaligned_loop_lb(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: scf.for
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_misaligned_loop_lb(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
+ %lb = arith.constant 1 : index
+ %ub = arith.constant 25 : index
+ %step = arith.constant 8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %r = scf.for %iv = %lb to %ub step %step iter_args(%out = %dst) -> tensor<32xf32> {
+ %slice = tensor.extract_slice %flat[%iv] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ %insert = tensor.insert_slice %slice into %out[%iv] [8] [1] : tensor<8xf32> into tensor<32xf32>
+ scf.yield %insert : tensor<32xf32>
+ }
+ return %r : tensor<32xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_collapse_misaligned_loop_step(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: scf.for
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_misaligned_loop_step(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
+ %lb = arith.constant 0 : index
+ %ub = arith.constant 25 : index
+ %step = arith.constant 4 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %r = scf.for %iv = %lb to %ub step %step iter_args(%out = %dst) -> tensor<32xf32> {
+ %slice = tensor.extract_slice %flat[%iv] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ %insert = tensor.insert_slice %slice into %out[%iv] [8] [1] : tensor<8xf32> into tensor<32xf32>
+ scf.yield %insert : tensor<32xf32>
+ }
+ return %r : tensor<32xf32>
+}
+
+// The data-flow fixed point also proves alignment of a loop-carried offset.
+// CHECK-LABEL: func.func @bubble_collapse_iter_arg(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: scf.for
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_iter_arg(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %r:2 = scf.for %iv = %c0 to %c4 step %c1 iter_args(%offset = %c0, %out = %dst) -> (index, tensor<32xf32>) {
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ %insert = tensor.insert_slice %slice into %out[%offset] [8] [1] : tensor<8xf32> into tensor<32xf32>
+ %next = arith.addi %offset, %c8 : index
+ scf.yield %next, %insert : index, tensor<32xf32>
+ }
+ return %r#1 : tensor<32xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_negative_multiplier(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_negative_multiplier(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %cneg8 = arith.constant -8 : index
+ %offset = arith.muli %row, %cneg8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_collapse_divisor_product_overflow(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_divisor_product_overflow(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c = arith.constant -9223372036854775808 : index
+ %c2 = arith.constant 2 : index
+ %a = arith.muli %row, %c : index
+ %offset = arith.muli %a, %c2 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// Loop-carried facts must be weakened when the recurrence loses alignment.
+// CHECK-LABEL: func.func @no_bubble_collapse_variant_iter_arg(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: scf.for
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_variant_iter_arg(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %r:2 = scf.for %iv = %c0 to %c4 step %c1 iter_args(%offset = %c0, %out = %dst) -> (index, tensor<32xf32>) {
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ %insert = tensor.insert_slice %slice into %out[%offset] [8] [1] : tensor<8xf32> into tensor<32xf32>
+ %next = arith.addi %offset, %c1 : index
+ scf.yield %next, %insert : index, tensor<32xf32>
+ }
+ return %r#1 : tensor<32xf32>
+}
+
+// A general divisor cannot be propagated through division after wraparound.
+// CHECK-LABEL: func.func @no_bubble_collapse_wrapping_division(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_wrapping_division(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
+ %c24 = arith.constant 24 : index
+ %c3 = arith.constant 3 : index
+ %base = arith.muli %row, %c24 : index
+ %offset = arith.divui %base, %c3 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice : tensor<8xf32>
+}
+
+// Analysis of the enclosing scope must tolerate unrelated integer operations.
+// CHECK-LABEL: func.func @bubble_collapse_with_unrelated_integer_ops(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_with_unrelated_integer_ops(%src: tensor<4x8xf32>, %row: index, %wide: i128) -> (tensor<8xf32>, i128, index) {
+ %c = arith.constant 170141183460469231731687303715884105728 : i128
+ %wide_product = arith.muli %wide, %c : i128
+ %zero = arith.constant 0 : index
+ %undefined = arith.divui %row, %zero : index
+ %c8 = arith.constant 8 : index
+ %offset = arith.muli %row, %c8 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ return %slice, %wide_product, %undefined : tensor<8xf32>, i128, index
+}
+
+// The first batch cannot prove alignment. Replacing the loop creates a new IV
+// whose step is eight; the second batch must analyze it rather than reuse facts.
+// CHECK-LABEL: func.func @bubble_collapse_replaced_loop(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[C8:.*]] = arith.constant 8 : index
+// CHECK: scf.for %[[IV:.*]] = {{.*}} step %[[C8]]
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 8] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_collapse_replaced_loop(%src: tensor<4x8xf32>, %out: tensor<32xf32>) -> tensor<32xf32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c32 = arith.constant 32 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %r = scf.for %iv = %c0 to %c32 step %c1 iter_args(%acc = %out) -> tensor<32xf32> {
+ %slice = tensor.extract_slice %flat[%iv] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ %next = tensor.insert_slice %slice into %acc[%iv] [8] [1] : tensor<8xf32> into tensor<32xf32>
+ scf.yield %next : tensor<32xf32>
+ } {test.replacement_step = 8 : i64}
+ return %r : tensor<32xf32>
+}
+
+// Refreshing analysis must not invent alignment for a misaligned lower bound.
+// CHECK-LABEL: func.func @no_bubble_collapse_replaced_loop_misaligned_lb(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: scf.for
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_collapse_replaced_loop_misaligned_lb(%src: tensor<4x8xf32>, %out: tensor<32xf32>) -> tensor<32xf32> {
+ %c1 = arith.constant 1 : index
+ %c32 = arith.constant 32 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
+ %r = scf.for %iv = %c1 to %c32 step %c1 iter_args(%acc = %out) -> tensor<32xf32> {
+ %slice = tensor.extract_slice %flat[%iv] [8] [1] : tensor<32xf32> to tensor<8xf32>
+ %next = tensor.insert_slice %slice into %acc[%iv] [8] [1] : tensor<8xf32> into tensor<32xf32>
+ scf.yield %next : tensor<32xf32>
+ } {test.replacement_step = 8 : i64}
+ return %r : tensor<32xf32>
+}
diff --git a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
index 3490782e188d6a..31232ee7151b2b 100644
--- a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
+++ b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
@@ -492,283 +492,3 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
-
-// -----
-
-// Power-of-two alignment through integer expressions and loop bounds.
-
-// CHECK-LABEL: func.func @bubble_collapse_muli_offset(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
-// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
-// CHECK-SAME: [1, 8] [1, 1]
-// CHECK: tensor.collapse_shape %[[SLICE]]
-func.func @bubble_collapse_muli_offset(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
- %c8 = arith.constant 8 : index
- %offset = arith.muli %row, %c8 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
- return %slice : tensor<8xf32>
-}
-
-// CHECK-LABEL: func.func @bubble_collapse_addi_offset(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
-// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
-// CHECK-SAME: [1, 8] [1, 1]
-// CHECK: tensor.collapse_shape %[[SLICE]]
-func.func @bubble_collapse_addi_offset(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
- %c8 = arith.constant 8 : index
- %base = arith.muli %row, %c8 : index
- %offset = arith.addi %base, %c8 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
- return %slice : tensor<8xf32>
-}
-
-// CHECK-LABEL: func.func @bubble_collapse_subi_offset(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
-// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
-// CHECK-SAME: [1, 8] [1, 1]
-// CHECK: tensor.collapse_shape %[[SLICE]]
-func.func @bubble_collapse_subi_offset(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
- %c8 = arith.constant 8 : index
- %base = arith.muli %row, %c8 : index
- %offset = arith.subi %base, %c8 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
- return %slice : tensor<8xf32>
-}
-
-// CHECK-LABEL: func.func @bubble_collapse_select_offset(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
-// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
-// CHECK-SAME: [1, 8] [1, 1]
-// CHECK: tensor.collapse_shape %[[SLICE]]
-func.func @bubble_collapse_select_offset(%src: tensor<4x8xf32>, %row: index, %cond: i1) -> tensor<8xf32> {
- %c8 = arith.constant 8 : index
- %a = arith.muli %row, %c8 : index
- %b = arith.addi %a, %c8 : index
- %offset = arith.select %cond, %a, %b : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
- return %slice : tensor<8xf32>
-}
-
-// CHECK-LABEL: func.func @no_bubble_collapse_misaligned_add(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
-// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
-// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_collapse_misaligned_add(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
- %c8 = arith.constant 8 : index
- %c1 = arith.constant 1 : index
- %base = arith.muli %row, %c8 : index
- %offset = arith.addi %base, %c1 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
- return %slice : tensor<8xf32>
-}
-
-// CHECK-LABEL: func.func @no_bubble_collapse_misaligned_select(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
-// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
-// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_collapse_misaligned_select(%src: tensor<4x8xf32>, %row: index, %cond: i1) -> tensor<8xf32> {
- %c8 = arith.constant 8 : index
- %c1 = arith.constant 1 : index
- %a = arith.muli %row, %c8 : index
- %b = arith.addi %a, %c1 : index
- %offset = arith.select %cond, %a, %b : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
- return %slice : tensor<8xf32>
-}
-
-// CHECK-LABEL: func.func @no_bubble_collapse_non_power_of_two_muli(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
-// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
-// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_collapse_non_power_of_two_muli(%src: tensor<4x6xf32>, %row: index) -> tensor<6xf32> {
- %c6 = arith.constant 6 : index
- %offset = arith.muli %row, %c6 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
- %slice = tensor.extract_slice %flat[%offset] [6] [1] : tensor<24xf32> to tensor<6xf32>
- return %slice : tensor<6xf32>
-}
-
-// CHECK-LABEL: func.func @no_bubble_collapse_unsupported_division(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
-// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
-// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_collapse_unsupported_division(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
- %c8 = arith.constant 8 : index
- %c2 = arith.constant 2 : index
- %base = arith.muli %row, %c8 : index
- %offset = arith.divui %base, %c2 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
- return %slice : tensor<8xf32>
-}
-
-// CHECK-LABEL: func.func @no_bubble_collapse_query_budget(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
-// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
-// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_collapse_query_budget(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
- %c8 = arith.constant 8 : index
- %v0 = arith.muli %row, %c8 : index
- %v1 = arith.addi %v0, %v0 : index
- %v2 = arith.addi %v1, %v1 : index
- %v3 = arith.addi %v2, %v2 : index
- %v4 = arith.addi %v3, %v3 : index
- %v5 = arith.addi %v4, %v4 : index
- %v6 = arith.addi %v5, %v5 : index
- %offset = arith.addi %v6, %v6 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
- return %slice : tensor<8xf32>
-}
-
-// CHECK-LABEL: func.func @bubble_collapse_multiple_groups(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8x2x4xf32>
-// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
-// CHECK-SAME: [1, 8, 1, 4] [1, 1, 1, 1]
-// CHECK: tensor.collapse_shape %[[SLICE]]
-func.func @bubble_collapse_multiple_groups(%src: tensor<4x8x2x4xf32>, %row: index, %col: index) -> tensor<8x4xf32> {
- %c8 = arith.constant 8 : index
- %c4 = arith.constant 4 : index
- %a = arith.muli %row, %c8 : index
- %b = arith.muli %col, %c4 : index
- %flat = tensor.collapse_shape %src [[0, 1], [2, 3]] : tensor<4x8x2x4xf32> into tensor<32x8xf32>
- %slice = tensor.extract_slice %flat[%a, %b] [8, 4] [1, 1] : tensor<32x8xf32> to tensor<8x4xf32>
- return %slice : tensor<8x4xf32>
-}
-
-// CHECK-LABEL: func.func @bubble_collapse_loop_iv(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
-// CHECK: scf.for
-// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
-// CHECK-SAME: [1, 8] [1, 1]
-// CHECK: tensor.collapse_shape %[[SLICE]]
-func.func @bubble_collapse_loop_iv(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
- %lb = arith.constant 0 : index
- %ub = arith.constant 25 : index
- %step = arith.constant 8 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %r = scf.for %iv = %lb to %ub step %step iter_args(%out = %dst) -> tensor<32xf32> {
- %slice = tensor.extract_slice %flat[%iv] [8] [1] : tensor<32xf32> to tensor<8xf32>
- %insert = tensor.insert_slice %slice into %out[%iv] [8] [1] : tensor<8xf32> into tensor<32xf32>
- scf.yield %insert : tensor<32xf32>
- }
- return %r : tensor<32xf32>
-}
-
-// CHECK-LABEL: func.func @bubble_collapse_aligned_loop_lb(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
-// CHECK: scf.for
-// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
-// CHECK-SAME: [1, 8] [1, 1]
-// CHECK: tensor.collapse_shape %[[SLICE]]
-func.func @bubble_collapse_aligned_loop_lb(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
- %lb = arith.constant 8 : index
- %ub = arith.constant 25 : index
- %step = arith.constant 8 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %r = scf.for %iv = %lb to %ub step %step iter_args(%out = %dst) -> tensor<32xf32> {
- %slice = tensor.extract_slice %flat[%iv] [8] [1] : tensor<32xf32> to tensor<8xf32>
- %insert = tensor.insert_slice %slice into %out[%iv] [8] [1] : tensor<8xf32> into tensor<32xf32>
- scf.yield %insert : tensor<32xf32>
- }
- return %r : tensor<32xf32>
-}
-
-// CHECK-LABEL: func.func @no_bubble_collapse_misaligned_loop_lb(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
-// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
-// CHECK: scf.for
-// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_collapse_misaligned_loop_lb(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
- %lb = arith.constant 1 : index
- %ub = arith.constant 25 : index
- %step = arith.constant 8 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %r = scf.for %iv = %lb to %ub step %step iter_args(%out = %dst) -> tensor<32xf32> {
- %slice = tensor.extract_slice %flat[%iv] [8] [1] : tensor<32xf32> to tensor<8xf32>
- %insert = tensor.insert_slice %slice into %out[%iv] [8] [1] : tensor<8xf32> into tensor<32xf32>
- scf.yield %insert : tensor<32xf32>
- }
- return %r : tensor<32xf32>
-}
-
-// CHECK-LABEL: func.func @no_bubble_collapse_misaligned_loop_step(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
-// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
-// CHECK: scf.for
-// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_collapse_misaligned_loop_step(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
- %lb = arith.constant 0 : index
- %ub = arith.constant 25 : index
- %step = arith.constant 4 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %r = scf.for %iv = %lb to %ub step %step iter_args(%out = %dst) -> tensor<32xf32> {
- %slice = tensor.extract_slice %flat[%iv] [8] [1] : tensor<32xf32> to tensor<8xf32>
- %insert = tensor.insert_slice %slice into %out[%iv] [8] [1] : tensor<8xf32> into tensor<32xf32>
- scf.yield %insert : tensor<32xf32>
- }
- return %r : tensor<32xf32>
-}
-
-// CHECK-LABEL: func.func @no_bubble_collapse_iter_arg(
-// CHECK: %[[FLAT:.*]] = tensor.collapse_shape
-// CHECK: scf.for
-// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_collapse_iter_arg(%src: tensor<4x8xf32>, %dst: tensor<32xf32>) -> tensor<32xf32> {
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
- %c4 = arith.constant 4 : index
- %c8 = arith.constant 8 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %r:2 = scf.for %iv = %c0 to %c4 step %c1 iter_args(%offset = %c0, %out = %dst) -> (index, tensor<32xf32>) {
- %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
- %insert = tensor.insert_slice %slice into %out[%offset] [8] [1] : tensor<8xf32> into tensor<32xf32>
- %next = arith.addi %offset, %c8 : index
- scf.yield %next, %insert : index, tensor<32xf32>
- }
- return %r#1 : tensor<32xf32>
-}
-
-// CHECK-LABEL: func.func @bubble_collapse_negative_multiplier(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
-// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
-// CHECK-SAME: [1, 8] [1, 1]
-// CHECK: tensor.collapse_shape %[[SLICE]]
-func.func @bubble_collapse_negative_multiplier(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
- %cneg8 = arith.constant -8 : index
- %offset = arith.muli %row, %cneg8 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
- return %slice : tensor<8xf32>
-}
-
-// CHECK-LABEL: func.func @bubble_collapse_divisor_product_overflow(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
-// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
-// CHECK-SAME: [1, 8] [1, 1]
-// CHECK: tensor.collapse_shape %[[SLICE]]
-func.func @bubble_collapse_divisor_product_overflow(%src: tensor<4x8xf32>, %row: index) -> tensor<8xf32> {
- %c = arith.constant -9223372036854775808 : index
- %c2 = arith.constant 2 : index
- %a = arith.muli %row, %c : index
- %offset = arith.muli %a, %c2 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x8xf32> into tensor<32xf32>
- %slice = tensor.extract_slice %flat[%offset] [8] [1] : tensor<32xf32> to tensor<8xf32>
- return %slice : tensor<8xf32>
-}
-
-module attributes {transform.with_named_sequence} {
- transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
- %func = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
- transform.apply_patterns to %func {
- transform.apply_patterns.tensor.bubble_up_extract_slice
- } : !transform.any_op
- transform.yield
- }
-}
diff --git a/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp b/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
index 2fb12a2f7f1980..cc327242342a96 100644
--- a/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
@@ -77,6 +77,11 @@ struct TestTensorTransforms
llvm::cl::desc("Test folding of expand_shape/collapse_shape"),
llvm::cl::init(false)};
+ Option<bool> testBubbleUpExtractSliceAlignment{
+ *this, "test-bubble-up-extract-slice-alignment",
+ llvm::cl::desc("Test analysis-backed extract_slice bubbling"),
+ llvm::cl::init(false)};
+
Option<bool> testFoldExtractFromCollapseShape{
*this, "test-fold-extract-from-collapse-shape",
llvm::cl::desc("Test folding of extract from collapse_shape"),
@@ -377,6 +382,35 @@ static LogicalResult testTrackingListenerReplacements(Operation *rootOp) {
void TestTensorTransforms::runOnOperation() {
Operation *rootOp = getOperation();
+ if (testBubbleUpExtractSliceAlignment) {
+ IRRewriter rewriter(rootOp->getContext());
+ auto bubbleSlices = [&]() {
+ SmallVector<Operation *> slices;
+ rootOp->walk(
+ [&](tensor::ExtractSliceOp slice) { slices.push_back(slice); });
+ return tensor::bubbleUpExtractSlicesThroughCollapseShape(rewriter, rootOp,
+ slices);
+ };
+ if (failed(bubbleSlices()))
+ return signalPassFailure();
+ SmallVector<scf::ForOp> loopsToReplace;
+ rootOp->walk([&](scf::ForOp loop) {
+ if (loop->hasAttr("test.replacement_step"))
+ loopsToReplace.push_back(loop);
+ });
+ for (scf::ForOp loop : loopsToReplace) {
+ rewriter.setInsertionPoint(loop);
+ auto stepAttr = loop->getAttrOfType<IntegerAttr>("test.replacement_step");
+ Value step = arith::ConstantIndexOp::create(rewriter, loop.getLoc(),
+ stepAttr.getInt());
+ auto replacement = cast<scf::ForOp>(rewriter.clone(*loop));
+ replacement.getStepMutable().assign(step);
+ replacement->removeAttr("test.replacement_step");
+ rewriter.replaceOp(loop, replacement);
+ }
+ if (!loopsToReplace.empty() && failed(bubbleSlices()))
+ return signalPassFailure();
+ }
if (testFoldConstantExtractSlice)
applyFoldConstantExtractSlicePatterns(rootOp);
if (testFoldConsecutiveInsertExtractSlice)
More information about the Mlir-commits
mailing list