[Mlir-commits] [mlir] [mlir][tensor] Recognize SCF IV divisibility when bubbling up extract_slice (PR #236185)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Oct 31 03:27:04 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/5] [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/5] [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)
>From aa211b9c7c1c63766a7d42277fec4d28ae1aac2f Mon Sep 17 00:00:00 2001
From: meinie0826 <929449726 at qq.com>
Date: Wed, 28 Oct 2026 01:35:24 +0800
Subject: [PATCH 3/5] [mlir][tensor] Separate divisibility analysis from slice
rewrite planning
---
.../Arith/Analysis/IntegerDivisibility.h | 46 ++++
.../SCF/Transforms/TileUsingInterface.h | 15 +-
.../Dialect/Tensor/Transforms/Transforms.h | 17 +-
.../lib/Dialect/Arith/Analysis/CMakeLists.txt | 12 +
.../Arith/Analysis/IntegerDivisibility.cpp | 96 +++++++
mlir/lib/Dialect/Arith/CMakeLists.txt | 1 +
.../TransformOps/LinalgTransformOps.cpp | 30 +--
.../SCF/Transforms/TileUsingInterface.cpp | 17 +-
.../Dialect/Tensor/Transforms/CMakeLists.txt | 5 +-
.../Tensor/Transforms/ReshapePatterns.cpp | 251 +++++++-----------
.../integer-power-of-two-divisibility.mlir | 83 ++++++
mlir/test/lib/Analysis/CMakeLists.txt | 1 +
.../TestIntegerDivisibilityAnalysis.cpp | 16 +-
13 files changed, 396 insertions(+), 194 deletions(-)
create mode 100644 mlir/include/mlir/Dialect/Arith/Analysis/IntegerDivisibility.h
create mode 100644 mlir/lib/Dialect/Arith/Analysis/CMakeLists.txt
create mode 100644 mlir/lib/Dialect/Arith/Analysis/IntegerDivisibility.cpp
create mode 100644 mlir/test/Analysis/DataFlow/integer-power-of-two-divisibility.mlir
diff --git a/mlir/include/mlir/Dialect/Arith/Analysis/IntegerDivisibility.h b/mlir/include/mlir/Dialect/Arith/Analysis/IntegerDivisibility.h
new file mode 100644
index 00000000000000..378d08e8ec5b41
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Arith/Analysis/IntegerDivisibility.h
@@ -0,0 +1,46 @@
+//===- IntegerDivisibility.h - Divisibility analysis -*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ARITH_ANALYSIS_INTEGERDIVISIBILITY_H
+#define MLIR_DIALECT_ARITH_ANALYSIS_INTEGERDIVISIBILITY_H
+
+#include "mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h"
+
+namespace mlir::arith {
+/// Infer power-of-two divisibility using the existing divisibility lattice and
+/// fixed-point/control-flow analysis. Support arith integer constants,
+/// add/subtract/multiply/select, and loop induction variables with bounds
+/// representable in int64_t. Other operation transfer functions are
+/// pessimistic.
+///
+/// Restrict facts to powers of two so modular arithmetic cannot invalidate
+/// them. Nonzero divisor products are saturated rather than encoded as zero:
+/// a zero divisor is reserved for a value known to be zero. Wide constants are
+/// supported without narrowing or taking signed absolute values.
+///
+/// Like IntegerDivisibilityAnalysis, this requires DeadCodeAnalysis. Load
+/// SparseConstantPropagation to analyze loop induction variables. Results must
+/// not be reused after modifying the analyzed IR.
+class PowerOfTwoDivisibilityAnalysis
+ : public dataflow::IntegerDivisibilityAnalysis {
+public:
+ using IntegerDivisibilityAnalysis::IntegerDivisibilityAnalysis;
+
+ LogicalResult visitOperation(
+ Operation *op,
+ ArrayRef<const dataflow::IntegerDivisibilityLattice *> operands,
+ ArrayRef<dataflow::IntegerDivisibilityLattice *> results) override;
+
+ void visitNonControlFlowArguments(
+ Operation *op, const RegionSuccessor &successor,
+ ValueRange successorInputs,
+ ArrayRef<dataflow::IntegerDivisibilityLattice *> argLattices) override;
+};
+} // namespace mlir::arith
+
+#endif // MLIR_DIALECT_ARITH_ANALYSIS_INTEGERDIVISIBILITY_H
diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
index 10a616428c2619..de34d52bf75e66 100644
--- a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
+++ b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
@@ -317,12 +317,15 @@ struct SCFTileAndFuseOptions {
/// `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;
+ /// Optional batch transformation before the greedy cleanup patterns, called
+ /// after tiling and after each fusion step. Return the surviving and inserted
+ /// operations to pass to the patterns. Use the supplied rewriter so that the
+ /// extract_slice worklist tracks replacements. This provides a phase boundary
+ /// for analysis-backed transformations without retaining analysis across
+ /// greedy rewrites.
+ using PreCleanupFnTy = std::function<FailureOr<SmallVector<Operation *>>(
+ RewriterBase &, ArrayRef<Operation *>)>;
+ PreCleanupFnTy preCleanupFn;
};
/// 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 552ea7074d24cc..cf9fdaa815332e 100644
--- a/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
@@ -171,15 +171,14 @@ getCollapsedExtractSliceInfo(OpBuilder &b, tensor::ExtractSliceOp sliceOp,
/// the shape of `expandedValue`.
///
/// This fails when the specified expansion cannot be represented by a valid
-/// 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);
+/// ExtractSliceOp.
+LogicalResult
+getExpandedExtractSliceInfo(OpBuilder &b, tensor::ExtractSliceOp sliceOp,
+ ArrayRef<ReassociationIndices> reassociation,
+ Value expandedValue,
+ SmallVectorImpl<OpFoldResult> &expandedOffsets,
+ SmallVectorImpl<OpFoldResult> &expandedSizes,
+ SmallVectorImpl<OpFoldResult> &expandedStrides);
} // namespace tensor
} // namespace mlir
diff --git a/mlir/lib/Dialect/Arith/Analysis/CMakeLists.txt b/mlir/lib/Dialect/Arith/Analysis/CMakeLists.txt
new file mode 100644
index 00000000000000..f3885185a92bf4
--- /dev/null
+++ b/mlir/lib/Dialect/Arith/Analysis/CMakeLists.txt
@@ -0,0 +1,12 @@
+add_mlir_dialect_library(MLIRArithAnalysis
+ IntegerDivisibility.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arith/Analysis
+
+ LINK_LIBS PUBLIC
+ MLIRAnalysisDataFlow
+ MLIRArithDialect
+ MLIRIR
+ MLIRLoopLikeInterface
+ )
diff --git a/mlir/lib/Dialect/Arith/Analysis/IntegerDivisibility.cpp b/mlir/lib/Dialect/Arith/Analysis/IntegerDivisibility.cpp
new file mode 100644
index 00000000000000..02ee9599051a57
--- /dev/null
+++ b/mlir/lib/Dialect/Arith/Analysis/IntegerDivisibility.cpp
@@ -0,0 +1,96 @@
+//===- IntegerDivisibility.cpp - Divisibility analysis -------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/Analysis/IntegerDivisibility.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "llvm/ADT/STLExtras.h"
+
+#include <numeric>
+
+using namespace mlir;
+using namespace mlir::arith;
+
+LogicalResult PowerOfTwoDivisibilityAnalysis::visitOperation(
+ Operation *op,
+ ArrayRef<const dataflow::IntegerDivisibilityLattice *> operands,
+ ArrayRef<dataflow::IntegerDivisibilityLattice *> results) {
+ if (auto constant = dyn_cast<ConstantOp>(op)) {
+ auto attr = dyn_cast<IntegerAttr>(constant.getValue());
+ if (!attr) {
+ setAllToEntryStates(results);
+ return success();
+ }
+ const APInt &value = attr.getValue();
+ 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();
+ }
+ // Division does not preserve the modular-arithmetic proof, even if a
+ // general (non-power-of-two) divisor of its numerator is known.
+ if (!isa<AddIOp, SubIOp, MulIOp, SelectOp>(op)) {
+ setAllToEntryStates(results);
+ return success();
+ }
+ if (llvm::any_of(operands, [](const auto *lattice) {
+ return lattice->getValue().isUninitialized();
+ }))
+ return success();
+
+ if (isa<MulIOp>(op)) {
+ uint64_t lhs = operands[0]->getValue().getValue().udiv();
+ uint64_t rhs = operands[1]->getValue().getValue().udiv();
+ if (lhs && rhs && lhs > UINT64_MAX / rhs) {
+ // Each nonzero operand divisor is a power of two. A product outside
+ // uint64_t is still divisible by 2^63, but need not be exactly zero.
+ uint64_t divisor = uint64_t{1} << 63;
+ propagateIfChanged(
+ results.front(),
+ results.front()->join(ConstantIntDivisibility(divisor, divisor)));
+ return success();
+ }
+ }
+ return IntegerDivisibilityAnalysis::visitOperation(op, operands, results);
+}
+
+void PowerOfTwoDivisibilityAnalysis::visitNonControlFlowArguments(
+ Operation *op, const RegionSuccessor &successor, ValueRange successorInputs,
+ ArrayRef<dataflow::IntegerDivisibilityLattice *> argLattices) {
+ auto loop = dyn_cast<LoopLikeOpInterface>(op);
+ auto ivs =
+ loop ? loop.getLoopInductionVars() : std::optional<SmallVector<Value>>();
+ // The existing loop-bound query represents constants in int64_t.
+ 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);
+ if (!ivs)
+ return;
+ // Reuse the loop-bound/step inference, retaining only its modular facts.
+ for (Value iv : *ivs) {
+ auto *lattice = getLatticeElement(iv);
+ if (lattice->getValue().isUninitialized())
+ continue;
+ const auto &div = lattice->getValue().getValue();
+ uint64_t udiv = div.udiv() & (uint64_t{0} - div.udiv());
+ uint64_t sdiv = div.sdiv() & (uint64_t{0} - div.sdiv());
+ // Both interpretations have the same power-of-two factors. Joining also
+ // discards odd factors before any arithmetic transfer consumes the IV.
+ uint64_t divisor = std::gcd(udiv, sdiv);
+ propagateIfChanged(
+ lattice, lattice->join(ConstantIntDivisibility(divisor, divisor)));
+ }
+}
diff --git a/mlir/lib/Dialect/Arith/CMakeLists.txt b/mlir/lib/Dialect/Arith/CMakeLists.txt
index 31167e6af908b9..35c4201f21454b 100644
--- a/mlir/lib/Dialect/Arith/CMakeLists.txt
+++ b/mlir/lib/Dialect/Arith/CMakeLists.txt
@@ -1,3 +1,4 @@
+add_subdirectory(Analysis)
add_subdirectory(IR)
add_subdirectory(Transforms)
add_subdirectory(Utils)
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index b82f37ac59a9f4..a65ae371767522 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -753,26 +753,16 @@ transform::FuseOp::apply(transform::TransformRewriter &rewriter,
tensor::ExtractSliceOp::getCanonicalizationPatterns(patterns, context);
tensor::populateMergeConsecutiveInsertExtractSlicePatterns(patterns);
tensor::populateBubbleUpExtractSliceOpPatterns(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));
- };
+ tileAndFuseOptions.cleanupPatterns = std::move(patterns);
+ tileAndFuseOptions.preCleanupFn = [](RewriterBase &rewriter,
+ ArrayRef<Operation *> ops) {
+ Operation *scope = ops.front();
+ while (!scope->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
+ scope->getParentOp())
+ scope = scope->getParentOp();
+ return tensor::bubbleUpExtractSlicesThroughCollapseShape(rewriter, scope,
+ ops);
+ };
}
size_t numLoops;
diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index 57b7428098de76..57f0af06a70e2f 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -1649,14 +1649,13 @@ 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;
+ scf::SCFTileAndFuseOptions::PreCleanupFnTy preCleanupFn;
};
SliceTrackingListener::SliceTrackingListener(
const scf::SCFTileAndFuseOptions &options) {
patterns = options.cleanupPatterns;
- cleanupFn = options.cleanupFn;
+ preCleanupFn = options.preCleanupFn;
}
LogicalResult
@@ -1668,13 +1667,19 @@ SliceTrackingListener::insertAndApplyPatterns(ArrayRef<Operation *> ops) {
if (ops.empty())
return success();
- if (cleanupFn)
- return cleanupFn(ops, this);
+ SmallVector<Operation *> cleanedOps(ops);
+ if (preCleanupFn) {
+ IRRewriter rewriter(ops.front()->getContext(), this);
+ auto result = preCleanupFn(rewriter, ops);
+ if (failed(result))
+ return failure();
+ cleanedOps = std::move(*result);
+ }
if (!patterns)
return success();
return applyOpPatternsGreedily(
- ops, patterns.value(),
+ cleanedOps, patterns.value(),
GreedyRewriteConfig().setListener(this).setStrictness(
GreedyRewriteStrictness::ExistingAndNewOps));
}
diff --git a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
index 9198e307017a14..3b19c4c9b49a09 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
@@ -20,18 +20,17 @@ add_mlir_dialect_library(MLIRTensorTransforms
MLIRTensorTransformsIncGen
LINK_LIBS PUBLIC
- MLIRAnalysis
+ MLIRAnalysisDataFlow
MLIRAffineDialect
MLIRAffineTransforms
MLIRAffineUtils
+ MLIRArithAnalysis
MLIRArithDialect
MLIRArithUtils
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 6f2a3705392723..1210d9e3f75999 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
@@ -10,14 +10,12 @@
#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/Analysis/IntegerDivisibility.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/ADT/SmallPtrSet.h"
@@ -626,10 +624,11 @@ LogicalResult mlir::tensor::getCollapsedExtractSliceInfo(
return success();
}
-// 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) {
+// Check static and affine divisibility, or consult an additional proof. The
+// proof must remain valid until all slice information is computed.
+static bool
+isMultipleOf(OpFoldResult ofr, int64_t factor,
+ llvm::function_ref<bool(Value, int64_t)> isKnownMultipleOf) {
if (factor <= 0)
return false;
if (auto constant = getConstantIntValue(ofr))
@@ -638,7 +637,7 @@ static bool isMultipleOf(OpFoldResult ofr, int64_t factor,
Value value = dyn_cast<Value>(ofr);
if (!value)
return false;
- if (isAligned && isAligned(value, factor))
+ if (isKnownMultipleOf && isKnownMultipleOf(value, factor))
return true;
auto applyOp = value.getDefiningOp<affine::AffineApplyOp>();
if (!applyOp)
@@ -662,7 +661,7 @@ static LogicalResult computeExpandedSliceInfoForReassocGroup(
OpBuilder &b, OpFoldResult collapsedSize, OpFoldResult collapsedOffset,
const ReassociationIndices &reassocIndices, ArrayRef<int64_t> expandedShape,
SmallVectorImpl<OpFoldResult> &groupSizes,
- llvm::function_ref<bool(Value, int64_t)> isAligned) {
+ llvm::function_ref<bool(Value, int64_t)> isKnownMultipleOf) {
assert(groupSizes.empty() && "Group sizes must be empty");
// The first case is when there's only one non-unit dimension in the
// reassociation group.
@@ -744,7 +743,7 @@ static LogicalResult computeExpandedSliceInfoForReassocGroup(
// Check dynamic/static offset divisibility.
currentOffsetDivisor *= expandedShapeSize;
- if (!isMultipleOf(collapsedOffset, currentOffsetDivisor, isAligned))
+ if (!isMultipleOf(collapsedOffset, currentOffsetDivisor, isKnownMultipleOf))
return failure();
// Trailing dims get full shape and zero offset.
@@ -772,7 +771,7 @@ static LogicalResult computeExpandedSliceInfoForReassocGroup(
// to check the validity of the range.
if ((expandedShapeSize % currentCollapsedsize) != 0)
return failure();
- if (!isMultipleOf(collapsedOffset, staticSize.value(), isAligned))
+ if (!isMultipleOf(collapsedOffset, staticSize.value(), isKnownMultipleOf))
return failure();
}
// Slicing dimension gets the remaining collapsed size.
@@ -794,130 +793,73 @@ static LogicalResult computeExpandedSliceInfoForReassocGroup(
return success();
}
-LogicalResult mlir::tensor::getExpandedExtractSliceInfo(
- OpBuilder &b, tensor::ExtractSliceOp sliceOp,
+/// Validate the geometry and compute sizes without inserting operations.
+static LogicalResult computeExpandedExtractSliceSizes(
+ OpBuilder &b, ExtractSliceOp sliceOp,
ArrayRef<ReassociationIndices> reassociation, Value expandedValue,
- SmallVectorImpl<OpFoldResult> &expandedOffsets,
SmallVectorImpl<OpFoldResult> &expandedSizes,
- SmallVectorImpl<OpFoldResult> &expandedStrides,
- llvm::function_ref<bool(Value, int64_t)> isAligned) {
- if (!sliceOp.hasUnitStride()) {
+ llvm::function_ref<bool(Value, int64_t)> isKnownMultipleOf) {
+ if (!sliceOp.hasUnitStride())
return failure();
- }
-
- // The tensor.extract_slice before applying the pattern works on the result
- // of the tensor.collapse_shape, so variables (i.e. inputs for
- // ExtractSliceOp) referring to the state before applying the pattern are
- // named with the prefix "collapsed", and ones referring to the state after
- // applying the pattern are named with the prefix "expanded".
SmallVector<OpFoldResult> collapsedOffsets = sliceOp.getMixedOffsets();
SmallVector<OpFoldResult> collapsedSizes = sliceOp.getMixedSizes();
if (static_cast<size_t>(sliceOp.getResultType().getRank()) !=
- collapsedSizes.size()) {
+ collapsedSizes.size())
return failure();
- }
- // Compute new offsets, sizes, and strides for tensor.extract_slice.
- // The new tensor.extract_slice will work on a tensor that has has a rank
- // equal to the rank of the src of the collapse_shape. In each iteration of
- // the loop, the offsets and sizes will be computed per reassociation group.
ArrayRef<int64_t> expandedShape =
cast<RankedTensorType>(expandedValue.getType()).getShape();
- SmallVector<SmallVector<OpFoldResult>> groupResults;
+ SmallVector<OpFoldResult> sizes;
for (auto [collapsedSize, collapsedOffset, reassocIndices] :
llvm::zip_equal(collapsedSizes, collapsedOffsets, reassociation)) {
-
SmallVector<OpFoldResult> groupSizes;
- LogicalResult result = computeExpandedSliceInfoForReassocGroup(
- b, collapsedSize, collapsedOffset, reassocIndices, expandedShape,
- groupSizes, isAligned);
- if (failed(result))
+ if (failed(computeExpandedSliceInfoForReassocGroup(
+ b, collapsedSize, collapsedOffset, reassocIndices, expandedShape,
+ groupSizes, isKnownMultipleOf)))
return failure();
- groupResults.emplace_back(groupSizes);
+ sizes.append(groupSizes);
}
+ expandedSizes.append(sizes);
+ return success();
+}
- expandedStrides.resize(expandedShape.size(), b.getIndexAttr(1));
+/// Materialize offsets and strides after the slice geometry has been checked.
+static void materializeExpandedExtractSliceInfo(
+ OpBuilder &b, ExtractSliceOp sliceOp,
+ ArrayRef<ReassociationIndices> reassociation, Value expandedValue,
+ SmallVectorImpl<OpFoldResult> &expandedOffsets,
+ SmallVectorImpl<OpFoldResult> &expandedStrides) {
+ auto expandedType = cast<RankedTensorType>(expandedValue.getType());
+ expandedStrides.resize(expandedType.getRank(), b.getIndexAttr(1));
+ SmallVector<OpFoldResult> collapsedOffsets = sliceOp.getMixedOffsets();
for (auto [groupIdx, reassocIndices] : llvm::enumerate(reassociation)) {
- auto &sizes = groupResults[groupIdx];
- expandedSizes.append(sizes);
-
SmallVector<OpFoldResult> basis;
for (int64_t expandedShapeIdx : reassocIndices)
basis.push_back(tensor::getMixedSize(b, sliceOp.getLoc(), expandedValue,
expandedShapeIdx));
- OpFoldResult collapsedOffset = collapsedOffsets[groupIdx];
- Value offsetVal =
- getValueOrCreateConstantIndexOp(b, sliceOp.getLoc(), collapsedOffset);
+ Value offsetVal = getValueOrCreateConstantIndexOp(
+ b, sliceOp.getLoc(), collapsedOffsets[groupIdx]);
auto delinearizeOp = affine::AffineDelinearizeIndexOp::create(
b, sliceOp.getLoc(), offsetVal, basis, /*hasOuterBound=*/true);
for (OpResult result : delinearizeOp.getResults())
expandedOffsets.push_back(result);
}
- 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
+LogicalResult mlir::tensor::getExpandedExtractSliceInfo(
+ OpBuilder &b, tensor::ExtractSliceOp sliceOp,
+ ArrayRef<ReassociationIndices> reassociation, Value expandedValue,
+ SmallVectorImpl<OpFoldResult> &expandedOffsets,
+ SmallVectorImpl<OpFoldResult> &expandedSizes,
+ SmallVectorImpl<OpFoldResult> &expandedStrides) {
+ if (failed(computeExpandedExtractSliceSizes(
+ b, sliceOp, reassociation, expandedValue, expandedSizes, nullptr)))
+ return failure();
+ materializeExpandedExtractSliceInfo(b, sliceOp, reassociation, expandedValue,
+ expandedOffsets, expandedStrides);
+ return success();
+}
FailureOr<SmallVector<Operation *>>
mlir::tensor::bubbleUpExtractSlicesThroughCollapseShape(
@@ -929,63 +871,74 @@ mlir::tensor::bubbleUpExtractSlicesThroughCollapseShape(
}))
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;
+ struct SliceRewritePlan {
+ Operation *op;
+ CollapseShapeOp collapse;
+ SmallVector<OpFoldResult> sizes;
};
+ SmallVector<SliceRewritePlan> plans;
+ {
+ DataFlowSolver solver;
+ solver.load<dataflow::DeadCodeAnalysis>();
+ solver.load<dataflow::SparseConstantPropagation>();
+ solver.load<arith::PowerOfTwoDivisibilityAnalysis>();
+ if (failed(solver.initializeAndRun(analysisScope)))
+ return failure();
- // 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;
+ auto isKnownMultipleOf = [&](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;
+ };
+
+ // Plan the batch without inserting or modifying any operation. Only sizes
+ // and handles to existing operations survive the analysis scope.
+ 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();
+ plans.push_back({op, collapse, {}});
+ if (collapse &&
+ failed(computeExpandedExtractSliceSizes(
+ rewriter, slice, collapse.getReassociationIndices(),
+ collapse.getSrc(), plans.back().sizes, isKnownMultipleOf)))
+ plans.back().collapse = nullptr;
+ }
+ }
+
+ // The solver has been destroyed: rewriting cannot consume stale facts.
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);
+ for (const SliceRewritePlan &plan : plans) {
+ if (!plan.collapse) {
+ cleanedOperations.push_back(plan.op);
continue;
}
+ auto slice = cast<ExtractSliceOp>(plan.op);
+ auto collapse = plan.collapse;
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);
+ SmallVector<OpFoldResult> offsets, strides;
+ materializeExpandedExtractSliceInfo(rewriter, slice,
+ collapse.getReassociationIndices(),
+ collapse.getSrc(), offsets, strides);
+ auto expandedSlice =
+ ExtractSliceOp::create(rewriter, slice.getLoc(), collapse.getSrc(),
+ offsets, plan.sizes, strides);
auto replacement = CollapseShapeOp::create(
rewriter, slice.getLoc(), slice.getType(), expandedSlice,
collapse.getReassociationIndices());
- replacements.emplace_back(slice, replacement);
+ rewriter.replaceOp(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;
}
diff --git a/mlir/test/Analysis/DataFlow/integer-power-of-two-divisibility.mlir b/mlir/test/Analysis/DataFlow/integer-power-of-two-divisibility.mlir
new file mode 100644
index 00000000000000..44d0c741f0a556
--- /dev/null
+++ b/mlir/test/Analysis/DataFlow/integer-power-of-two-divisibility.mlir
@@ -0,0 +1,83 @@
+// RUN: mlir-opt %s -test-int-divisibility-analysis="power-of-two" | FileCheck %s
+
+// CHECK-LABEL: @arithmetic
+func.func @arithmetic(%a: index, %b: index, %p: i1) {
+ %c8 = arith.constant 8 : index
+ %c12 = arith.constant 12 : index
+ %a8 = arith.muli %a, %c8 : index
+ %b12 = arith.muli %b, %c12 : index
+ %sum = arith.addi %a8, %b12 : index
+ // CHECK: divisibility = "udiv = 4, sdiv = 4"
+ "test.int_divisibility"(%sum) : (index) -> ()
+ %difference = arith.subi %a8, %b12 : index
+ // CHECK: divisibility = "udiv = 4, sdiv = 4"
+ "test.int_divisibility"(%difference) : (index) -> ()
+ %selected = arith.select %p, %a8, %b12 : index
+ // CHECK: divisibility = "udiv = 4, sdiv = 4"
+ "test.int_divisibility"(%selected) : (index) -> ()
+ return
+}
+
+// CHECK-LABEL: @wide_constants_and_products
+func.func @wide_constants_and_products(%a: i128) {
+ %wide = arith.constant 170141183460469231731687303715884105728 : i128
+ // CHECK: divisibility = "udiv = 9223372036854775808, sdiv = 9223372036854775808"
+ "test.int_divisibility"(%wide) : (i128) -> ()
+ %a_wide = arith.muli %a, %wide : i128
+ %two = arith.constant 2 : i128
+ %product = arith.muli %a_wide, %two : i128
+ // Saturated divisor products must not assert that the value is exactly zero.
+ // CHECK: divisibility = "udiv = 9223372036854775808, sdiv = 9223372036854775808"
+ "test.int_divisibility"(%product) : (i128) -> ()
+ %zero = arith.constant 0 : i128
+ %zero_product = arith.muli %a, %zero : i128
+ // CHECK: divisibility = "udiv = 0, sdiv = 0"
+ "test.int_divisibility"(%zero_product) : (i128) -> ()
+ return
+}
+
+// CHECK-LABEL: @division_is_pessimistic
+func.func @division_is_pessimistic(%a: index) {
+ %c24 = arith.constant 24 : index
+ %c3 = arith.constant 3 : index
+ %product = arith.muli %a, %c24 : index
+ %quotient = arith.divui %product, %c3 : index
+ // CHECK: divisibility = "udiv = 1, sdiv = 1"
+ "test.int_divisibility"(%quotient) : (index) -> ()
+ %zero = arith.constant 0 : index
+ %undefined = arith.divui %a, %zero : index
+ // CHECK: divisibility = "udiv = 1, sdiv = 1"
+ "test.int_divisibility"(%undefined) : (index) -> ()
+ return
+}
+
+// CHECK-LABEL: @induction_variables
+func.func @induction_variables() {
+ %c0 = arith.constant 0 : index
+ %c6 = arith.constant 6 : index
+ %c24 = arith.constant 24 : index
+ scf.for %iv = %c0 to %c24 step %c6 {
+ // Retain the power-of-two part of gcd(0, 6), not the odd factor three.
+ // CHECK: divisibility = "udiv = 2, sdiv = 2"
+ "test.int_divisibility"(%iv) : (index) -> ()
+ }
+ %c3 = arith.constant 3 : index
+ scf.for %iv = %c3 to %c24 step %c6 {
+ // CHECK: divisibility = "udiv = 1, sdiv = 1"
+ "test.int_divisibility"(%iv) : (index) -> ()
+ }
+ return
+}
+
+// CHECK-LABEL: @wide_induction_variable
+func.func @wide_induction_variable() {
+ %c0 = arith.constant 0 : i128
+ %c8 = arith.constant 8 : i128
+ %c32 = arith.constant 32 : i128
+ scf.for %iv = %c0 to %c32 step %c8 : i128 {
+ // Loop-bound inference representing constants in int64_t is not invoked.
+ // CHECK: divisibility = "udiv = 1, sdiv = 1"
+ "test.int_divisibility"(%iv) : (i128) -> ()
+ }
+ return
+}
diff --git a/mlir/test/lib/Analysis/CMakeLists.txt b/mlir/test/lib/Analysis/CMakeLists.txt
index 706c18eb2da5eb..5a9a3c5cfceec7 100644
--- a/mlir/test/lib/Analysis/CMakeLists.txt
+++ b/mlir/test/lib/Analysis/CMakeLists.txt
@@ -49,6 +49,7 @@ mlir_target_link_libraries(MLIRTestAnalysis PUBLIC
MLIRAnalysisBase
MLIRAnalysisDataFlow
MLIRArithDialect
+ MLIRArithAnalysis
MLIRFunctionInterfaces
MLIRMemRefDialect
MLIRPass
diff --git a/mlir/test/lib/Analysis/DataFlow/TestIntegerDivisibilityAnalysis.cpp b/mlir/test/lib/Analysis/DataFlow/TestIntegerDivisibilityAnalysis.cpp
index 6f7b0b3c89bde2..6d6109e6e377d7 100644
--- a/mlir/test/lib/Analysis/DataFlow/TestIntegerDivisibilityAnalysis.cpp
+++ b/mlir/test/lib/Analysis/DataFlow/TestIntegerDivisibilityAnalysis.cpp
@@ -11,6 +11,7 @@
#include "mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h"
#include "mlir/Analysis/DataFlowFramework.h"
#include "mlir/Dialect/Affine/IR/AffineDialect.h"
+#include "mlir/Dialect/Arith/Analysis/IntegerDivisibility.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/DialectRegistry.h"
@@ -29,6 +30,16 @@ struct TestIntegerDivisibilityAnalysisPass
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
TestIntegerDivisibilityAnalysisPass)
+ TestIntegerDivisibilityAnalysisPass() = default;
+ TestIntegerDivisibilityAnalysisPass(
+ const TestIntegerDivisibilityAnalysisPass &pass)
+ : PassWrapper(pass) {}
+
+ Option<bool> powerOfTwo{
+ *this, "power-of-two",
+ llvm::cl::desc("Infer only modular power-of-two divisibility"),
+ llvm::cl::init(false)};
+
StringRef getArgument() const override {
return "test-int-divisibility-analysis";
}
@@ -63,7 +74,10 @@ struct TestIntegerDivisibilityAnalysisPass
// visitNonControlFlowArguments and analyze arguments like loop induction
// variables.
solver.load<SparseConstantPropagation>();
- solver.load<IntegerDivisibilityAnalysis>();
+ if (powerOfTwo)
+ solver.load<arith::PowerOfTwoDivisibilityAnalysis>();
+ else
+ solver.load<IntegerDivisibilityAnalysis>();
if (failed(solver.initializeAndRun(rootOp)))
return signalPassFailure();
>From f2792ea3e04fe1d3bf294b336d9125e531a4f77e Mon Sep 17 00:00:00 2001
From: meinie0826 <929449726 at qq.com>
Date: Wed, 28 Oct 2026 01:56:28 +0800
Subject: [PATCH 4/5] [mlir][tensor] Prove reshape slice divisibility from
static SCF IV bounds
Use the existing reshape pattern and loop interface to prove direct SCF induction variables divisible by a factor from constant lower bounds and steps. SCF no-overflow semantics makes this valid for arbitrary positive factors. Keep unsupported values conservative and avoid dataflow analysis or tiling hooks.
---
.../Arith/Analysis/IntegerDivisibility.h | 46 --
.../SCF/Transforms/TileUsingInterface.h | 10 -
.../Dialect/Tensor/Transforms/Transforms.h | 11 +-
.../lib/Dialect/Arith/Analysis/CMakeLists.txt | 12 -
.../Arith/Analysis/IntegerDivisibility.cpp | 96 ----
mlir/lib/Dialect/Arith/CMakeLists.txt | 1 -
.../TransformOps/LinalgTransformOps.cpp | 9 -
.../SCF/Transforms/TileUsingInterface.cpp | 24 +-
.../Dialect/Tensor/Transforms/CMakeLists.txt | 3 +-
.../Tensor/Transforms/ReshapePatterns.cpp | 213 +++------
.../integer-power-of-two-divisibility.mlir | 83 ----
.../Dialect/Linalg/transform-op-fuse.mlir | 44 +-
.../bubble-up-extract-slice-alignment.mlir | 431 +++++-------------
mlir/test/lib/Analysis/CMakeLists.txt | 1 -
.../TestIntegerDivisibilityAnalysis.cpp | 16 +-
.../Dialect/Tensor/TestTensorTransforms.cpp | 34 --
16 files changed, 225 insertions(+), 809 deletions(-)
delete mode 100644 mlir/include/mlir/Dialect/Arith/Analysis/IntegerDivisibility.h
delete mode 100644 mlir/lib/Dialect/Arith/Analysis/CMakeLists.txt
delete mode 100644 mlir/lib/Dialect/Arith/Analysis/IntegerDivisibility.cpp
delete mode 100644 mlir/test/Analysis/DataFlow/integer-power-of-two-divisibility.mlir
diff --git a/mlir/include/mlir/Dialect/Arith/Analysis/IntegerDivisibility.h b/mlir/include/mlir/Dialect/Arith/Analysis/IntegerDivisibility.h
deleted file mode 100644
index 378d08e8ec5b41..00000000000000
--- a/mlir/include/mlir/Dialect/Arith/Analysis/IntegerDivisibility.h
+++ /dev/null
@@ -1,46 +0,0 @@
-//===- IntegerDivisibility.h - Divisibility analysis -*- C++ -*-===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_DIALECT_ARITH_ANALYSIS_INTEGERDIVISIBILITY_H
-#define MLIR_DIALECT_ARITH_ANALYSIS_INTEGERDIVISIBILITY_H
-
-#include "mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h"
-
-namespace mlir::arith {
-/// Infer power-of-two divisibility using the existing divisibility lattice and
-/// fixed-point/control-flow analysis. Support arith integer constants,
-/// add/subtract/multiply/select, and loop induction variables with bounds
-/// representable in int64_t. Other operation transfer functions are
-/// pessimistic.
-///
-/// Restrict facts to powers of two so modular arithmetic cannot invalidate
-/// them. Nonzero divisor products are saturated rather than encoded as zero:
-/// a zero divisor is reserved for a value known to be zero. Wide constants are
-/// supported without narrowing or taking signed absolute values.
-///
-/// Like IntegerDivisibilityAnalysis, this requires DeadCodeAnalysis. Load
-/// SparseConstantPropagation to analyze loop induction variables. Results must
-/// not be reused after modifying the analyzed IR.
-class PowerOfTwoDivisibilityAnalysis
- : public dataflow::IntegerDivisibilityAnalysis {
-public:
- using IntegerDivisibilityAnalysis::IntegerDivisibilityAnalysis;
-
- LogicalResult visitOperation(
- Operation *op,
- ArrayRef<const dataflow::IntegerDivisibilityLattice *> operands,
- ArrayRef<dataflow::IntegerDivisibilityLattice *> results) override;
-
- void visitNonControlFlowArguments(
- Operation *op, const RegionSuccessor &successor,
- ValueRange successorInputs,
- ArrayRef<dataflow::IntegerDivisibilityLattice *> argLattices) override;
-};
-} // namespace mlir::arith
-
-#endif // MLIR_DIALECT_ARITH_ANALYSIS_INTEGERDIVISIBILITY_H
diff --git a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
index de34d52bf75e66..f457da8b3d71df 100644
--- a/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
+++ b/mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h
@@ -316,16 +316,6 @@ 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 batch transformation before the greedy cleanup patterns, called
- /// after tiling and after each fusion step. Return the surviving and inserted
- /// operations to pass to the patterns. Use the supplied rewriter so that the
- /// extract_slice worklist tracks replacements. This provides a phase boundary
- /// for analysis-backed transformations without retaining analysis across
- /// greedy rewrites.
- using PreCleanupFnTy = std::function<FailureOr<SmallVector<Operation *>>(
- RewriterBase &, ArrayRef<Operation *>)>;
- PreCleanupFnTy preCleanupFn;
};
/// 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 cf9fdaa815332e..341b3b1dc3472e 100644
--- a/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
@@ -70,16 +70,6 @@ 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(
@@ -179,6 +169,7 @@ getExpandedExtractSliceInfo(OpBuilder &b, tensor::ExtractSliceOp sliceOp,
SmallVectorImpl<OpFoldResult> &expandedOffsets,
SmallVectorImpl<OpFoldResult> &expandedSizes,
SmallVectorImpl<OpFoldResult> &expandedStrides);
+
} // namespace tensor
} // namespace mlir
diff --git a/mlir/lib/Dialect/Arith/Analysis/CMakeLists.txt b/mlir/lib/Dialect/Arith/Analysis/CMakeLists.txt
deleted file mode 100644
index f3885185a92bf4..00000000000000
--- a/mlir/lib/Dialect/Arith/Analysis/CMakeLists.txt
+++ /dev/null
@@ -1,12 +0,0 @@
-add_mlir_dialect_library(MLIRArithAnalysis
- IntegerDivisibility.cpp
-
- ADDITIONAL_HEADER_DIRS
- ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arith/Analysis
-
- LINK_LIBS PUBLIC
- MLIRAnalysisDataFlow
- MLIRArithDialect
- MLIRIR
- MLIRLoopLikeInterface
- )
diff --git a/mlir/lib/Dialect/Arith/Analysis/IntegerDivisibility.cpp b/mlir/lib/Dialect/Arith/Analysis/IntegerDivisibility.cpp
deleted file mode 100644
index 02ee9599051a57..00000000000000
--- a/mlir/lib/Dialect/Arith/Analysis/IntegerDivisibility.cpp
+++ /dev/null
@@ -1,96 +0,0 @@
-//===- IntegerDivisibility.cpp - Divisibility analysis -------------------===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Dialect/Arith/Analysis/IntegerDivisibility.h"
-#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Interfaces/LoopLikeInterface.h"
-#include "llvm/ADT/STLExtras.h"
-
-#include <numeric>
-
-using namespace mlir;
-using namespace mlir::arith;
-
-LogicalResult PowerOfTwoDivisibilityAnalysis::visitOperation(
- Operation *op,
- ArrayRef<const dataflow::IntegerDivisibilityLattice *> operands,
- ArrayRef<dataflow::IntegerDivisibilityLattice *> results) {
- if (auto constant = dyn_cast<ConstantOp>(op)) {
- auto attr = dyn_cast<IntegerAttr>(constant.getValue());
- if (!attr) {
- setAllToEntryStates(results);
- return success();
- }
- const APInt &value = attr.getValue();
- 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();
- }
- // Division does not preserve the modular-arithmetic proof, even if a
- // general (non-power-of-two) divisor of its numerator is known.
- if (!isa<AddIOp, SubIOp, MulIOp, SelectOp>(op)) {
- setAllToEntryStates(results);
- return success();
- }
- if (llvm::any_of(operands, [](const auto *lattice) {
- return lattice->getValue().isUninitialized();
- }))
- return success();
-
- if (isa<MulIOp>(op)) {
- uint64_t lhs = operands[0]->getValue().getValue().udiv();
- uint64_t rhs = operands[1]->getValue().getValue().udiv();
- if (lhs && rhs && lhs > UINT64_MAX / rhs) {
- // Each nonzero operand divisor is a power of two. A product outside
- // uint64_t is still divisible by 2^63, but need not be exactly zero.
- uint64_t divisor = uint64_t{1} << 63;
- propagateIfChanged(
- results.front(),
- results.front()->join(ConstantIntDivisibility(divisor, divisor)));
- return success();
- }
- }
- return IntegerDivisibilityAnalysis::visitOperation(op, operands, results);
-}
-
-void PowerOfTwoDivisibilityAnalysis::visitNonControlFlowArguments(
- Operation *op, const RegionSuccessor &successor, ValueRange successorInputs,
- ArrayRef<dataflow::IntegerDivisibilityLattice *> argLattices) {
- auto loop = dyn_cast<LoopLikeOpInterface>(op);
- auto ivs =
- loop ? loop.getLoopInductionVars() : std::optional<SmallVector<Value>>();
- // The existing loop-bound query represents constants in int64_t.
- 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);
- if (!ivs)
- return;
- // Reuse the loop-bound/step inference, retaining only its modular facts.
- for (Value iv : *ivs) {
- auto *lattice = getLatticeElement(iv);
- if (lattice->getValue().isUninitialized())
- continue;
- const auto &div = lattice->getValue().getValue();
- uint64_t udiv = div.udiv() & (uint64_t{0} - div.udiv());
- uint64_t sdiv = div.sdiv() & (uint64_t{0} - div.sdiv());
- // Both interpretations have the same power-of-two factors. Joining also
- // discards odd factors before any arithmetic transfer consumes the IV.
- uint64_t divisor = std::gcd(udiv, sdiv);
- propagateIfChanged(
- lattice, lattice->join(ConstantIntDivisibility(divisor, divisor)));
- }
-}
diff --git a/mlir/lib/Dialect/Arith/CMakeLists.txt b/mlir/lib/Dialect/Arith/CMakeLists.txt
index 35c4201f21454b..31167e6af908b9 100644
--- a/mlir/lib/Dialect/Arith/CMakeLists.txt
+++ b/mlir/lib/Dialect/Arith/CMakeLists.txt
@@ -1,4 +1,3 @@
-add_subdirectory(Analysis)
add_subdirectory(IR)
add_subdirectory(Transforms)
add_subdirectory(Utils)
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index a65ae371767522..816d7499845789 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -754,15 +754,6 @@ transform::FuseOp::apply(transform::TransformRewriter &rewriter,
tensor::populateMergeConsecutiveInsertExtractSlicePatterns(patterns);
tensor::populateBubbleUpExtractSliceOpPatterns(patterns);
tileAndFuseOptions.cleanupPatterns = std::move(patterns);
- tileAndFuseOptions.preCleanupFn = [](RewriterBase &rewriter,
- ArrayRef<Operation *> ops) {
- Operation *scope = ops.front();
- while (!scope->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
- scope->getParentOp())
- scope = scope->getParentOp();
- return tensor::bubbleUpExtractSlicesThroughCollapseShape(rewriter, scope,
- ops);
- };
}
size_t numLoops;
diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index 57f0af06a70e2f..a8efe6cdb4415e 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -1619,7 +1619,8 @@ namespace {
/// fusion algorithm to apply cleanup patterns in between fusion steps.
class SliceTrackingListener : public RewriterBase::Listener {
public:
- explicit SliceTrackingListener(const scf::SCFTileAndFuseOptions &options);
+ explicit SliceTrackingListener(
+ std::optional<FrozenRewritePatternSet> patterns);
SliceTrackingListener() = default;
/// Adds the given list of operations to the worklist, and if present,
@@ -1649,13 +1650,11 @@ class SliceTrackingListener : public RewriterBase::Listener {
/// Optional pattern set to apply when adding new operations to the
/// worklist.
std::optional<FrozenRewritePatternSet> patterns = std::nullopt;
- scf::SCFTileAndFuseOptions::PreCleanupFnTy preCleanupFn;
};
SliceTrackingListener::SliceTrackingListener(
- const scf::SCFTileAndFuseOptions &options) {
- patterns = options.cleanupPatterns;
- preCleanupFn = options.preCleanupFn;
+ std::optional<FrozenRewritePatternSet> p) {
+ patterns = std::move(p);
}
LogicalResult
@@ -1665,21 +1664,11 @@ SliceTrackingListener::insertAndApplyPatterns(ArrayRef<Operation *> ops) {
worklist.push_back(slice);
}
- if (ops.empty())
- return success();
- SmallVector<Operation *> cleanedOps(ops);
- if (preCleanupFn) {
- IRRewriter rewriter(ops.front()->getContext(), this);
- auto result = preCleanupFn(rewriter, ops);
- if (failed(result))
- return failure();
- cleanedOps = std::move(*result);
- }
if (!patterns)
return success();
return applyOpPatternsGreedily(
- cleanedOps, patterns.value(),
+ ops, patterns.value(),
GreedyRewriteConfig().setListener(this).setStrictness(
GreedyRewriteStrictness::ExistingAndNewOps));
}
@@ -1817,7 +1806,8 @@ mlir::scf::tileConsumerAndFuseProducersUsingSCF(
SCFTileAndFuseOptions::ControlFnResult controlFnResult;
};
- SliceTrackingListener sliceTracker = SliceTrackingListener(options);
+ SliceTrackingListener sliceTracker =
+ SliceTrackingListener(options.cleanupPatterns);
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 3b19c4c9b49a09..726839f3b61291 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
@@ -20,17 +20,16 @@ add_mlir_dialect_library(MLIRTensorTransforms
MLIRTensorTransformsIncGen
LINK_LIBS PUBLIC
- MLIRAnalysisDataFlow
MLIRAffineDialect
MLIRAffineTransforms
MLIRAffineUtils
- MLIRArithAnalysis
MLIRArithDialect
MLIRArithUtils
MLIRDialectUtils
MLIRFuncDialect # Dependency of scalarize-single-element-tensor-return
MLIRIR
MLIRLinalgDialect
+ MLIRLoopLikeInterface
MLIRMemRefDialect
MLIRPass
MLIRRuntimeVerifiableOpInterface
diff --git a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
index 1210d9e3f75999..844ba029b28c1f 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/ReshapePatterns.cpp
@@ -6,21 +6,17 @@
//
//===----------------------------------------------------------------------===//
-#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/Analysis/IntegerDivisibility.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/Dialect/SCF/IR/SCF.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/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"
using namespace mlir;
using namespace mlir::tensor;
@@ -624,20 +620,51 @@ LogicalResult mlir::tensor::getCollapsedExtractSliceInfo(
return success();
}
-// Check static and affine divisibility, or consult an additional proof. The
-// proof must remain valid until all slice information is computed.
-static bool
-isMultipleOf(OpFoldResult ofr, int64_t factor,
- llvm::function_ref<bool(Value, int64_t)> isKnownMultipleOf) {
+/// SCF for/parallel IVs are lb + j * step in arbitrary-precision arithmetic;
+/// overflow of the terminating increment is undefined. Constant lower bounds
+/// and steps therefore prove divisibility by any positive factor. The generic
+/// LoopLikeOpInterface does not guarantee this arithmetic semantics.
+static bool isLoopIVMultipleOf(Value value, int64_t factor) {
+ auto arg = dyn_cast<BlockArgument>(value);
+ if (!arg)
+ return false;
+ Operation *op = arg.getOwner()->getParentOp();
+ if (!isa<scf::ForOp, scf::ParallelOp>(op))
+ return false;
+ auto loop = cast<LoopLikeOpInterface>(op);
+ auto ivs = loop.getLoopInductionVars();
+ auto lbs = loop.getLoopLowerBounds();
+ auto steps = loop.getLoopSteps();
+ if (!ivs || !lbs || !steps)
+ return false;
+ for (auto [iv, lb, step] : llvm::zip_equal(*ivs, *lbs, *steps)) {
+ if (iv != value)
+ continue;
+ auto lowerBound = getConstantIntValue(lb);
+ auto stride = getConstantIntValue(step);
+ // Bounds that require the unsigned interpretation of a negative int64_t
+ // are left unproven rather than using signed remainder for that value.
+ if (!lowerBound || !stride || *stride <= 0 ||
+ (op->hasAttr("unsignedCmp") && *lowerBound < 0))
+ return false;
+ return *lowerBound % factor == 0 && *stride % factor == 0;
+ }
+ return false;
+}
+
+// Checks if the `ofr` is a multiple of the `factor`.
+// Handles constants, SCF induction variables, and affine.apply results.
+static bool isMultipleOf(OpFoldResult ofr, int64_t factor) {
if (factor <= 0)
return false;
- if (auto constant = getConstantIntValue(ofr))
- return *constant % factor == 0;
+ std::optional<int64_t> staticValue = getConstantIntValue(ofr);
+ if (staticValue.has_value())
+ return staticValue.value() % factor == 0;
Value value = dyn_cast<Value>(ofr);
if (!value)
return false;
- if (isKnownMultipleOf && isKnownMultipleOf(value, factor))
+ if (isLoopIVMultipleOf(value, factor))
return true;
auto applyOp = value.getDefiningOp<affine::AffineApplyOp>();
if (!applyOp)
@@ -660,8 +687,7 @@ 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,
- llvm::function_ref<bool(Value, int64_t)> isKnownMultipleOf) {
+ SmallVectorImpl<OpFoldResult> &groupSizes) {
assert(groupSizes.empty() && "Group sizes must be empty");
// The first case is when there's only one non-unit dimension in the
// reassociation group.
@@ -743,7 +769,7 @@ static LogicalResult computeExpandedSliceInfoForReassocGroup(
// Check dynamic/static offset divisibility.
currentOffsetDivisor *= expandedShapeSize;
- if (!isMultipleOf(collapsedOffset, currentOffsetDivisor, isKnownMultipleOf))
+ if (!isMultipleOf(collapsedOffset, currentOffsetDivisor))
return failure();
// Trailing dims get full shape and zero offset.
@@ -771,7 +797,7 @@ static LogicalResult computeExpandedSliceInfoForReassocGroup(
// to check the validity of the range.
if ((expandedShapeSize % currentCollapsedsize) != 0)
return failure();
- if (!isMultipleOf(collapsedOffset, staticSize.value(), isKnownMultipleOf))
+ if (!isMultipleOf(collapsedOffset, staticSize.value()))
return failure();
}
// Slicing dimension gets the remaining collapsed size.
@@ -793,155 +819,68 @@ static LogicalResult computeExpandedSliceInfoForReassocGroup(
return success();
}
-/// Validate the geometry and compute sizes without inserting operations.
-static LogicalResult computeExpandedExtractSliceSizes(
- OpBuilder &b, ExtractSliceOp sliceOp,
+LogicalResult mlir::tensor::getExpandedExtractSliceInfo(
+ OpBuilder &b, tensor::ExtractSliceOp sliceOp,
ArrayRef<ReassociationIndices> reassociation, Value expandedValue,
+ SmallVectorImpl<OpFoldResult> &expandedOffsets,
SmallVectorImpl<OpFoldResult> &expandedSizes,
- llvm::function_ref<bool(Value, int64_t)> isKnownMultipleOf) {
- if (!sliceOp.hasUnitStride())
+ SmallVectorImpl<OpFoldResult> &expandedStrides) {
+ if (!sliceOp.hasUnitStride()) {
return failure();
+ }
+
+ // The tensor.extract_slice before applying the pattern works on the result
+ // of the tensor.collapse_shape, so variables (i.e. inputs for
+ // ExtractSliceOp) referring to the state before applying the pattern are
+ // named with the prefix "collapsed", and ones referring to the state after
+ // applying the pattern are named with the prefix "expanded".
SmallVector<OpFoldResult> collapsedOffsets = sliceOp.getMixedOffsets();
SmallVector<OpFoldResult> collapsedSizes = sliceOp.getMixedSizes();
if (static_cast<size_t>(sliceOp.getResultType().getRank()) !=
- collapsedSizes.size())
+ collapsedSizes.size()) {
return failure();
+ }
+ // Compute new offsets, sizes, and strides for tensor.extract_slice.
+ // The new tensor.extract_slice will work on a tensor that has has a rank
+ // equal to the rank of the src of the collapse_shape. In each iteration of
+ // the loop, the offsets and sizes will be computed per reassociation group.
ArrayRef<int64_t> expandedShape =
cast<RankedTensorType>(expandedValue.getType()).getShape();
- SmallVector<OpFoldResult> sizes;
+ SmallVector<SmallVector<OpFoldResult>> groupResults;
for (auto [collapsedSize, collapsedOffset, reassocIndices] :
llvm::zip_equal(collapsedSizes, collapsedOffsets, reassociation)) {
+
SmallVector<OpFoldResult> groupSizes;
- if (failed(computeExpandedSliceInfoForReassocGroup(
- b, collapsedSize, collapsedOffset, reassocIndices, expandedShape,
- groupSizes, isKnownMultipleOf)))
+ LogicalResult result = computeExpandedSliceInfoForReassocGroup(
+ b, collapsedSize, collapsedOffset, reassocIndices, expandedShape,
+ groupSizes);
+ if (failed(result))
return failure();
- sizes.append(groupSizes);
+ groupResults.emplace_back(groupSizes);
}
- expandedSizes.append(sizes);
- return success();
-}
-/// Materialize offsets and strides after the slice geometry has been checked.
-static void materializeExpandedExtractSliceInfo(
- OpBuilder &b, ExtractSliceOp sliceOp,
- ArrayRef<ReassociationIndices> reassociation, Value expandedValue,
- SmallVectorImpl<OpFoldResult> &expandedOffsets,
- SmallVectorImpl<OpFoldResult> &expandedStrides) {
- auto expandedType = cast<RankedTensorType>(expandedValue.getType());
- expandedStrides.resize(expandedType.getRank(), b.getIndexAttr(1));
- SmallVector<OpFoldResult> collapsedOffsets = sliceOp.getMixedOffsets();
+ expandedStrides.resize(expandedShape.size(), b.getIndexAttr(1));
for (auto [groupIdx, reassocIndices] : llvm::enumerate(reassociation)) {
+ auto &sizes = groupResults[groupIdx];
+ expandedSizes.append(sizes);
+
SmallVector<OpFoldResult> basis;
for (int64_t expandedShapeIdx : reassocIndices)
basis.push_back(tensor::getMixedSize(b, sliceOp.getLoc(), expandedValue,
expandedShapeIdx));
- Value offsetVal = getValueOrCreateConstantIndexOp(
- b, sliceOp.getLoc(), collapsedOffsets[groupIdx]);
+ OpFoldResult collapsedOffset = collapsedOffsets[groupIdx];
+ Value offsetVal =
+ getValueOrCreateConstantIndexOp(b, sliceOp.getLoc(), collapsedOffset);
auto delinearizeOp = affine::AffineDelinearizeIndexOp::create(
b, sliceOp.getLoc(), offsetVal, basis, /*hasOuterBound=*/true);
for (OpResult result : delinearizeOp.getResults())
expandedOffsets.push_back(result);
}
-}
-
-LogicalResult mlir::tensor::getExpandedExtractSliceInfo(
- OpBuilder &b, tensor::ExtractSliceOp sliceOp,
- ArrayRef<ReassociationIndices> reassociation, Value expandedValue,
- SmallVectorImpl<OpFoldResult> &expandedOffsets,
- SmallVectorImpl<OpFoldResult> &expandedSizes,
- SmallVectorImpl<OpFoldResult> &expandedStrides) {
- if (failed(computeExpandedExtractSliceSizes(
- b, sliceOp, reassociation, expandedValue, expandedSizes, nullptr)))
- return failure();
- materializeExpandedExtractSliceInfo(b, sliceOp, reassociation, expandedValue,
- expandedOffsets, expandedStrides);
return success();
}
-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);
-
- struct SliceRewritePlan {
- Operation *op;
- CollapseShapeOp collapse;
- SmallVector<OpFoldResult> sizes;
- };
- SmallVector<SliceRewritePlan> plans;
- {
- DataFlowSolver solver;
- solver.load<dataflow::DeadCodeAnalysis>();
- solver.load<dataflow::SparseConstantPropagation>();
- solver.load<arith::PowerOfTwoDivisibilityAnalysis>();
- if (failed(solver.initializeAndRun(analysisScope)))
- return failure();
-
- auto isKnownMultipleOf = [&](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;
- };
-
- // Plan the batch without inserting or modifying any operation. Only sizes
- // and handles to existing operations survive the analysis scope.
- 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();
- plans.push_back({op, collapse, {}});
- if (collapse &&
- failed(computeExpandedExtractSliceSizes(
- rewriter, slice, collapse.getReassociationIndices(),
- collapse.getSrc(), plans.back().sizes, isKnownMultipleOf)))
- plans.back().collapse = nullptr;
- }
- }
-
- // The solver has been destroyed: rewriting cannot consume stale facts.
- SmallVector<Operation *> cleanedOperations;
- OpBuilder::InsertionGuard guard(rewriter);
- for (const SliceRewritePlan &plan : plans) {
- if (!plan.collapse) {
- cleanedOperations.push_back(plan.op);
- continue;
- }
- auto slice = cast<ExtractSliceOp>(plan.op);
- auto collapse = plan.collapse;
- rewriter.setInsertionPoint(slice);
- SmallVector<OpFoldResult> offsets, strides;
- materializeExpandedExtractSliceInfo(rewriter, slice,
- collapse.getReassociationIndices(),
- collapse.getSrc(), offsets, strides);
- auto expandedSlice =
- ExtractSliceOp::create(rewriter, slice.getLoc(), collapse.getSrc(),
- offsets, plan.sizes, strides);
- auto replacement = CollapseShapeOp::create(
- rewriter, slice.getLoc(), slice.getType(), expandedSlice,
- collapse.getReassociationIndices());
- rewriter.replaceOp(slice, replacement);
- // Include the expanded slice so the greedy cleanup can visit its producer.
- cleanedOperations.push_back(expandedSlice);
- cleanedOperations.push_back(replacement);
- }
- return cleanedOperations;
-}
-
void mlir::tensor::populateReassociativeReshapeFoldingPatterns(
RewritePatternSet &patterns) {
patterns.add<FoldExpandOfRankReducingExtract, FoldExtractSliceOfExpandShape,
diff --git a/mlir/test/Analysis/DataFlow/integer-power-of-two-divisibility.mlir b/mlir/test/Analysis/DataFlow/integer-power-of-two-divisibility.mlir
deleted file mode 100644
index 44d0c741f0a556..00000000000000
--- a/mlir/test/Analysis/DataFlow/integer-power-of-two-divisibility.mlir
+++ /dev/null
@@ -1,83 +0,0 @@
-// RUN: mlir-opt %s -test-int-divisibility-analysis="power-of-two" | FileCheck %s
-
-// CHECK-LABEL: @arithmetic
-func.func @arithmetic(%a: index, %b: index, %p: i1) {
- %c8 = arith.constant 8 : index
- %c12 = arith.constant 12 : index
- %a8 = arith.muli %a, %c8 : index
- %b12 = arith.muli %b, %c12 : index
- %sum = arith.addi %a8, %b12 : index
- // CHECK: divisibility = "udiv = 4, sdiv = 4"
- "test.int_divisibility"(%sum) : (index) -> ()
- %difference = arith.subi %a8, %b12 : index
- // CHECK: divisibility = "udiv = 4, sdiv = 4"
- "test.int_divisibility"(%difference) : (index) -> ()
- %selected = arith.select %p, %a8, %b12 : index
- // CHECK: divisibility = "udiv = 4, sdiv = 4"
- "test.int_divisibility"(%selected) : (index) -> ()
- return
-}
-
-// CHECK-LABEL: @wide_constants_and_products
-func.func @wide_constants_and_products(%a: i128) {
- %wide = arith.constant 170141183460469231731687303715884105728 : i128
- // CHECK: divisibility = "udiv = 9223372036854775808, sdiv = 9223372036854775808"
- "test.int_divisibility"(%wide) : (i128) -> ()
- %a_wide = arith.muli %a, %wide : i128
- %two = arith.constant 2 : i128
- %product = arith.muli %a_wide, %two : i128
- // Saturated divisor products must not assert that the value is exactly zero.
- // CHECK: divisibility = "udiv = 9223372036854775808, sdiv = 9223372036854775808"
- "test.int_divisibility"(%product) : (i128) -> ()
- %zero = arith.constant 0 : i128
- %zero_product = arith.muli %a, %zero : i128
- // CHECK: divisibility = "udiv = 0, sdiv = 0"
- "test.int_divisibility"(%zero_product) : (i128) -> ()
- return
-}
-
-// CHECK-LABEL: @division_is_pessimistic
-func.func @division_is_pessimistic(%a: index) {
- %c24 = arith.constant 24 : index
- %c3 = arith.constant 3 : index
- %product = arith.muli %a, %c24 : index
- %quotient = arith.divui %product, %c3 : index
- // CHECK: divisibility = "udiv = 1, sdiv = 1"
- "test.int_divisibility"(%quotient) : (index) -> ()
- %zero = arith.constant 0 : index
- %undefined = arith.divui %a, %zero : index
- // CHECK: divisibility = "udiv = 1, sdiv = 1"
- "test.int_divisibility"(%undefined) : (index) -> ()
- return
-}
-
-// CHECK-LABEL: @induction_variables
-func.func @induction_variables() {
- %c0 = arith.constant 0 : index
- %c6 = arith.constant 6 : index
- %c24 = arith.constant 24 : index
- scf.for %iv = %c0 to %c24 step %c6 {
- // Retain the power-of-two part of gcd(0, 6), not the odd factor three.
- // CHECK: divisibility = "udiv = 2, sdiv = 2"
- "test.int_divisibility"(%iv) : (index) -> ()
- }
- %c3 = arith.constant 3 : index
- scf.for %iv = %c3 to %c24 step %c6 {
- // CHECK: divisibility = "udiv = 1, sdiv = 1"
- "test.int_divisibility"(%iv) : (index) -> ()
- }
- return
-}
-
-// CHECK-LABEL: @wide_induction_variable
-func.func @wide_induction_variable() {
- %c0 = arith.constant 0 : i128
- %c8 = arith.constant 8 : i128
- %c32 = arith.constant 32 : i128
- scf.for %iv = %c0 to %c32 step %c8 : i128 {
- // Loop-bound inference representing constants in int64_t is not invoked.
- // CHECK: divisibility = "udiv = 1, sdiv = 1"
- "test.int_divisibility"(%iv) : (i128) -> ()
- }
- return
-}
diff --git a/mlir/test/Dialect/Linalg/transform-op-fuse.mlir b/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
index 0a5aa41a89caa8..2912e6fc1e357a 100644
--- a/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
@@ -906,33 +906,33 @@ module attributes {transform.with_named_sequence} {
// CHECK-NOT: linalg.matmul
// CHECK: scf.for
// CHECK: %[[INIT:.*]] = linalg.fill
-// CHECK-SAME: -> tensor<1x8xf32>
+// CHECK-SAME: -> tensor<1x6xf32>
// CHECK: %[[MM:.*]] = linalg.matmul
-// CHECK-SAME: outs(%[[INIT]] : tensor<1x8xf32>) -> tensor<1x8xf32>
+// CHECK-SAME: outs(%[[INIT]] : tensor<1x6xf32>) -> tensor<1x6xf32>
// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[MM]]
// CHECK: linalg.generic
-// CHECK-SAME: ins(%[[FLAT]] : tensor<8xf32>)
+// CHECK-SAME: ins(%[[FLAT]] : tensor<6xf32>)
// 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>
+func.func @fuse_matmul_through_collapse(%a: tensor<6x4xf32>, %b: tensor<4x6xf32>) -> tensor<36xf32> {
+ %zero = arith.constant 0.0 : f32
+ %empty = tensor.empty() : tensor<6x6xf32>
+ %init = linalg.fill ins(%zero : f32) outs(%empty : tensor<6x6xf32>) -> tensor<6x6xf32>
+ %mm = linalg.matmul ins(%a, %b : tensor<6x4xf32>, tensor<4x6xf32>) outs(%init : tensor<6x6xf32>) -> tensor<6x6xf32>
+ %flat = tensor.collapse_shape %mm [[0, 1]] : tensor<6x6xf32> into tensor<36xf32>
+ %out = tensor.empty() : tensor<36xf32>
+ %r = linalg.generic {candidate_root, indexing_maps = [affine_map<(i)->(i)>, affine_map<(i)->(i)>], iterator_types = ["parallel"]}
+ ins(%flat : tensor<36xf32>) outs(%out : tensor<36xf32>) {
+ ^bb0(%v: f32, %old: f32):
+ %v1 = arith.maximumf %v, %zero : f32
+ linalg.yield %v1 : f32
+ } -> tensor<36xf32>
+ return %r : tensor<36xf32>
}
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
- }
+ 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 [6] apply_cleanup : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ transform.yield
+ }
}
diff --git a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-alignment.mlir b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-alignment.mlir
index 1764ea3b3592d1..0a58a68505275b 100644
--- a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-alignment.mlir
+++ b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-alignment.mlir
@@ -1,367 +1,170 @@
-// RUN: mlir-opt %s -test-tensor-transform-patterns="test-bubble-up-extract-slice-alignment" | FileCheck %s
+// RUN: mlir-opt %s -transform-interpreter | 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>
+// Non-power-of-two factors are valid for SCF IVs, including dynamic upper bounds.
+// CHECK-LABEL: func.func @bubble_loop_iv(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
+// CHECK: scf.for
// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
-// CHECK-SAME: [1, 8] [1, 1]
+// CHECK-SAME: [1, 6] [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>
+func.func @bubble_loop_iv(%src: tensor<4x6xf32>, %out: tensor<24xf32>, %ub: index) -> tensor<24xf32> {
+ %c0 = arith.constant 0 : index
+ %c6 = arith.constant 6 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %r = scf.for %iv = %c0 to %ub step %c6 iter_args(%acc = %out) -> tensor<24xf32> {
+ %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
+ scf.yield %next : tensor<24xf32>
+ }
+ return %r : tensor<24xf32>
}
-// CHECK-LABEL: func.func @bubble_collapse_select_offset(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK-LABEL: func.func @bubble_nonzero_lb(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
+// CHECK: scf.for
// 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-SAME: [1, 6] [1, 1]
+func.func @bubble_nonzero_lb(%src: tensor<4x6xf32>, %out: tensor<24xf32>) -> tensor<24xf32> {
+ %c6 = arith.constant 6 : index
+ %c24 = arith.constant 24 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %r = scf.for %iv = %c6 to %c24 step %c6 iter_args(%acc = %out) -> tensor<24xf32> {
+ %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
+ scf.yield %next : tensor<24xf32>
+ }
+ return %r : tensor<24xf32>
}
-// CHECK-LABEL: func.func @no_bubble_collapse_non_power_of_two_muli(
+// CHECK-LABEL: func.func @no_bubble_misaligned_lb(
// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: scf.for
// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_collapse_non_power_of_two_muli(%src: tensor<4x6xf32>, %row: index) -> tensor<6xf32> {
+func.func @no_bubble_misaligned_lb(%src: tensor<4x6xf32>, %out: tensor<24xf32>) -> tensor<24xf32> {
+ %c1 = arith.constant 1 : index
%c6 = arith.constant 6 : index
- %offset = arith.muli %row, %c6 : index
+ %c19 = arith.constant 19 : 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>
+ %r = scf.for %iv = %c1 to %c19 step %c6 iter_args(%acc = %out) -> tensor<24xf32> {
+ %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
+ scf.yield %next : tensor<24xf32>
+ }
+ return %r : tensor<24xf32>
}
-// CHECK-LABEL: func.func @no_bubble_collapse_unsupported_division(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK-LABEL: func.func @no_bubble_misaligned_step(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: scf.for
// 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
+func.func @no_bubble_misaligned_step(%src: tensor<4x6xf32>, %out: tensor<24xf32>) -> tensor<24xf32> {
+ %c0 = arith.constant 0 : 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>
+ %c19 = arith.constant 19 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %r = scf.for %iv = %c0 to %c19 step %c4 iter_args(%acc = %out) -> tensor<24xf32> {
+ %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
+ scf.yield %next : tensor<24xf32>
}
- return %r : tensor<32xf32>
+ return %r : tensor<24xf32>
}
-// CHECK-LABEL: func.func @no_bubble_collapse_misaligned_loop_lb(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK-LABEL: func.func @no_bubble_dynamic_lb(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
// 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>
+func.func @no_bubble_dynamic_lb(%src: tensor<4x6xf32>, %out: tensor<24xf32>, %lb: index) -> tensor<24xf32> {
+ %c6 = arith.constant 6 : index
+ %c24 = arith.constant 24 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %r = scf.for %iv = %lb to %c24 step %c6 iter_args(%acc = %out) -> tensor<24xf32> {
+ %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
+ scf.yield %next : tensor<24xf32>
}
- return %r : tensor<32xf32>
+ return %r : tensor<24xf32>
}
-// CHECK-LABEL: func.func @no_bubble_collapse_misaligned_loop_step(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x8xf32>
+// CHECK-LABEL: func.func @no_bubble_dynamic_step(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
// 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> {
+func.func @no_bubble_dynamic_step(%src: tensor<4x6xf32>, %out: tensor<24xf32>, %step: index) -> tensor<24xf32> {
%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>
+ %c24 = arith.constant 24 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %r = scf.for %iv = %c0 to %c24 step %step iter_args(%acc = %out) -> tensor<24xf32> {
+ %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
+ scf.yield %next : tensor<24xf32>
}
- 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>
+ return %r : tensor<24xf32>
}
-// 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>
+// The recurrence happens to be aligned, but this rule only handles IVs.
+// CHECK-LABEL: func.func @no_bubble_iter_arg(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
// 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> {
+func.func @no_bubble_iter_arg(%src: tensor<4x6xf32>, %out: tensor<24xf32>) -> tensor<24xf32> {
%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>
+ %c6 = arith.constant 6 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %r:2 = scf.for %iv = %c0 to %c4 step %c1 iter_args(%offset = %c0, %acc = %out) -> (index, tensor<24xf32>) {
+ %tile = tensor.extract_slice %flat[%offset] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ %next = tensor.insert_slice %tile into %acc[%offset] [6] [1] : tensor<6xf32> into tensor<24xf32>
+ %next_offset = arith.addi %offset, %c6 : index
+ scf.yield %next_offset, %next : index, tensor<24xf32>
}
- return %r#1 : tensor<32xf32>
+ return %r#1 : tensor<24xf32>
}
-// 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-LABEL: func.func @no_bubble_arithmetic(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
// 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>
+func.func @no_bubble_arithmetic(%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>
+ %tile = tensor.extract_slice %flat[%offset] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ return %tile : tensor<6xf32>
}
-// 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
-}
+func.func private @consume(tensor<6x6xf32>)
-// 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-LABEL: func.func @bubble_parallel_ivs(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6x4x6xf32>
+// CHECK: scf.parallel
// 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> {
+// CHECK-SAME: [1, 6, 1, 6] [1, 1, 1, 1]
+func.func @bubble_parallel_ivs(%src: tensor<4x6x4x6xf32>) {
%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>
+ %c6 = arith.constant 6 : index
+ %c24 = arith.constant 24 : index
+ %flat = tensor.collapse_shape %src [[0, 1], [2, 3]] : tensor<4x6x4x6xf32> into tensor<24x24xf32>
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c24, %c24) step (%c6, %c6) {
+ %tile = tensor.extract_slice %flat[%i, %j] [6, 6] [1, 1] : tensor<24x24xf32> to tensor<6x6xf32>
+ func.call @consume(%tile) : (tensor<6x6xf32>) -> ()
+ scf.reduce
+ }
+ return
}
-// 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>
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+ %funcs = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %funcs {
+ transform.apply_patterns.tensor.bubble_up_extract_slice
+ } : !transform.any_op
+ transform.yield
+ }
}
diff --git a/mlir/test/lib/Analysis/CMakeLists.txt b/mlir/test/lib/Analysis/CMakeLists.txt
index 5a9a3c5cfceec7..706c18eb2da5eb 100644
--- a/mlir/test/lib/Analysis/CMakeLists.txt
+++ b/mlir/test/lib/Analysis/CMakeLists.txt
@@ -49,7 +49,6 @@ mlir_target_link_libraries(MLIRTestAnalysis PUBLIC
MLIRAnalysisBase
MLIRAnalysisDataFlow
MLIRArithDialect
- MLIRArithAnalysis
MLIRFunctionInterfaces
MLIRMemRefDialect
MLIRPass
diff --git a/mlir/test/lib/Analysis/DataFlow/TestIntegerDivisibilityAnalysis.cpp b/mlir/test/lib/Analysis/DataFlow/TestIntegerDivisibilityAnalysis.cpp
index 6d6109e6e377d7..6f7b0b3c89bde2 100644
--- a/mlir/test/lib/Analysis/DataFlow/TestIntegerDivisibilityAnalysis.cpp
+++ b/mlir/test/lib/Analysis/DataFlow/TestIntegerDivisibilityAnalysis.cpp
@@ -11,7 +11,6 @@
#include "mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h"
#include "mlir/Analysis/DataFlowFramework.h"
#include "mlir/Dialect/Affine/IR/AffineDialect.h"
-#include "mlir/Dialect/Arith/Analysis/IntegerDivisibility.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/DialectRegistry.h"
@@ -30,16 +29,6 @@ struct TestIntegerDivisibilityAnalysisPass
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
TestIntegerDivisibilityAnalysisPass)
- TestIntegerDivisibilityAnalysisPass() = default;
- TestIntegerDivisibilityAnalysisPass(
- const TestIntegerDivisibilityAnalysisPass &pass)
- : PassWrapper(pass) {}
-
- Option<bool> powerOfTwo{
- *this, "power-of-two",
- llvm::cl::desc("Infer only modular power-of-two divisibility"),
- llvm::cl::init(false)};
-
StringRef getArgument() const override {
return "test-int-divisibility-analysis";
}
@@ -74,10 +63,7 @@ struct TestIntegerDivisibilityAnalysisPass
// visitNonControlFlowArguments and analyze arguments like loop induction
// variables.
solver.load<SparseConstantPropagation>();
- if (powerOfTwo)
- solver.load<arith::PowerOfTwoDivisibilityAnalysis>();
- else
- solver.load<IntegerDivisibilityAnalysis>();
+ solver.load<IntegerDivisibilityAnalysis>();
if (failed(solver.initializeAndRun(rootOp)))
return signalPassFailure();
diff --git a/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp b/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
index cc327242342a96..2fb12a2f7f1980 100644
--- a/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
@@ -77,11 +77,6 @@ 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"),
@@ -382,35 +377,6 @@ 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)
>From 0bee3d05c5b710e07266947eec75104e4ee8fbd0 Mon Sep 17 00:00:00 2001
From: meinie0826 <929449726 at qq.com>
Date: Wed, 28 Oct 2026 10:12:10 +0800
Subject: [PATCH 5/5] [mlir][tensor] Keep loop IV slice tests in existing
reshape pattern suite
---
.../bubble-up-extract-slice-alignment.mlir | 170 ------------------
.../Tensor/bubble-up-extract-slice-op.mlir | 159 ++++++++++++++++
2 files changed, 159 insertions(+), 170 deletions(-)
delete mode 100644 mlir/test/Dialect/Tensor/bubble-up-extract-slice-alignment.mlir
diff --git a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-alignment.mlir b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-alignment.mlir
deleted file mode 100644
index 0a58a68505275b..00000000000000
--- a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-alignment.mlir
+++ /dev/null
@@ -1,170 +0,0 @@
-// RUN: mlir-opt %s -transform-interpreter | FileCheck %s
-
-// Non-power-of-two factors are valid for SCF IVs, including dynamic upper bounds.
-// CHECK-LABEL: func.func @bubble_loop_iv(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
-// CHECK: scf.for
-// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
-// CHECK-SAME: [1, 6] [1, 1]
-// CHECK: tensor.collapse_shape %[[SLICE]]
-func.func @bubble_loop_iv(%src: tensor<4x6xf32>, %out: tensor<24xf32>, %ub: index) -> tensor<24xf32> {
- %c0 = arith.constant 0 : index
- %c6 = arith.constant 6 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
- %r = scf.for %iv = %c0 to %ub step %c6 iter_args(%acc = %out) -> tensor<24xf32> {
- %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
- %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
- scf.yield %next : tensor<24xf32>
- }
- return %r : tensor<24xf32>
-}
-
-// CHECK-LABEL: func.func @bubble_nonzero_lb(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
-// CHECK: scf.for
-// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
-// CHECK-SAME: [1, 6] [1, 1]
-func.func @bubble_nonzero_lb(%src: tensor<4x6xf32>, %out: tensor<24xf32>) -> tensor<24xf32> {
- %c6 = arith.constant 6 : index
- %c24 = arith.constant 24 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
- %r = scf.for %iv = %c6 to %c24 step %c6 iter_args(%acc = %out) -> tensor<24xf32> {
- %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
- %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
- scf.yield %next : tensor<24xf32>
- }
- return %r : tensor<24xf32>
-}
-
-// CHECK-LABEL: func.func @no_bubble_misaligned_lb(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
-// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
-// CHECK: scf.for
-// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_misaligned_lb(%src: tensor<4x6xf32>, %out: tensor<24xf32>) -> tensor<24xf32> {
- %c1 = arith.constant 1 : index
- %c6 = arith.constant 6 : index
- %c19 = arith.constant 19 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
- %r = scf.for %iv = %c1 to %c19 step %c6 iter_args(%acc = %out) -> tensor<24xf32> {
- %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
- %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
- scf.yield %next : tensor<24xf32>
- }
- return %r : tensor<24xf32>
-}
-
-// CHECK-LABEL: func.func @no_bubble_misaligned_step(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
-// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
-// CHECK: scf.for
-// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_misaligned_step(%src: tensor<4x6xf32>, %out: tensor<24xf32>) -> tensor<24xf32> {
- %c0 = arith.constant 0 : index
- %c4 = arith.constant 4 : index
- %c19 = arith.constant 19 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
- %r = scf.for %iv = %c0 to %c19 step %c4 iter_args(%acc = %out) -> tensor<24xf32> {
- %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
- %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
- scf.yield %next : tensor<24xf32>
- }
- return %r : tensor<24xf32>
-}
-
-// CHECK-LABEL: func.func @no_bubble_dynamic_lb(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
-// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
-// CHECK: scf.for
-// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_dynamic_lb(%src: tensor<4x6xf32>, %out: tensor<24xf32>, %lb: index) -> tensor<24xf32> {
- %c6 = arith.constant 6 : index
- %c24 = arith.constant 24 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
- %r = scf.for %iv = %lb to %c24 step %c6 iter_args(%acc = %out) -> tensor<24xf32> {
- %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
- %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
- scf.yield %next : tensor<24xf32>
- }
- return %r : tensor<24xf32>
-}
-
-// CHECK-LABEL: func.func @no_bubble_dynamic_step(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
-// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
-// CHECK: scf.for
-// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_dynamic_step(%src: tensor<4x6xf32>, %out: tensor<24xf32>, %step: index) -> tensor<24xf32> {
- %c0 = arith.constant 0 : index
- %c24 = arith.constant 24 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
- %r = scf.for %iv = %c0 to %c24 step %step iter_args(%acc = %out) -> tensor<24xf32> {
- %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
- %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
- scf.yield %next : tensor<24xf32>
- }
- return %r : tensor<24xf32>
-}
-
-// The recurrence happens to be aligned, but this rule only handles IVs.
-// CHECK-LABEL: func.func @no_bubble_iter_arg(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
-// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
-// CHECK: scf.for
-// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_iter_arg(%src: tensor<4x6xf32>, %out: tensor<24xf32>) -> tensor<24xf32> {
- %c0 = arith.constant 0 : index
- %c1 = arith.constant 1 : index
- %c4 = arith.constant 4 : index
- %c6 = arith.constant 6 : index
- %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
- %r:2 = scf.for %iv = %c0 to %c4 step %c1 iter_args(%offset = %c0, %acc = %out) -> (index, tensor<24xf32>) {
- %tile = tensor.extract_slice %flat[%offset] [6] [1] : tensor<24xf32> to tensor<6xf32>
- %next = tensor.insert_slice %tile into %acc[%offset] [6] [1] : tensor<6xf32> into tensor<24xf32>
- %next_offset = arith.addi %offset, %c6 : index
- scf.yield %next_offset, %next : index, tensor<24xf32>
- }
- return %r#1 : tensor<24xf32>
-}
-
-// CHECK-LABEL: func.func @no_bubble_arithmetic(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
-// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
-// CHECK: tensor.extract_slice %[[FLAT]]
-func.func @no_bubble_arithmetic(%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>
- %tile = tensor.extract_slice %flat[%offset] [6] [1] : tensor<24xf32> to tensor<6xf32>
- return %tile : tensor<6xf32>
-}
-
-func.func private @consume(tensor<6x6xf32>)
-
-// CHECK-LABEL: func.func @bubble_parallel_ivs(
-// CHECK-SAME: %[[SRC:.*]]: tensor<4x6x4x6xf32>
-// CHECK: scf.parallel
-// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
-// CHECK-SAME: [1, 6, 1, 6] [1, 1, 1, 1]
-func.func @bubble_parallel_ivs(%src: tensor<4x6x4x6xf32>) {
- %c0 = arith.constant 0 : index
- %c6 = arith.constant 6 : index
- %c24 = arith.constant 24 : index
- %flat = tensor.collapse_shape %src [[0, 1], [2, 3]] : tensor<4x6x4x6xf32> into tensor<24x24xf32>
- scf.parallel (%i, %j) = (%c0, %c0) to (%c24, %c24) step (%c6, %c6) {
- %tile = tensor.extract_slice %flat[%i, %j] [6, 6] [1, 1] : tensor<24x24xf32> to tensor<6x6xf32>
- func.call @consume(%tile) : (tensor<6x6xf32>) -> ()
- scf.reduce
- }
- return
-}
-
-module attributes {transform.with_named_sequence} {
- transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
- %funcs = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.any_op
- transform.apply_patterns to %funcs {
- transform.apply_patterns.tensor.bubble_up_extract_slice
- } : !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..d0091f9032c193 100644
--- a/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
+++ b/mlir/test/Dialect/Tensor/bubble-up-extract-slice-op.mlir
@@ -483,6 +483,165 @@ func.func @no_bubble_up_extract_slice_affine_apply_dynamic_offset_multi_trailing
return %extracted_slice : tensor<24xf32>
}
+// Non-power-of-two factors are valid for SCF IVs, including dynamic upper bounds.
+// CHECK-LABEL: func.func @bubble_loop_iv(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
+// CHECK: scf.for
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 6] [1, 1]
+// CHECK: tensor.collapse_shape %[[SLICE]]
+func.func @bubble_loop_iv(%src: tensor<4x6xf32>, %out: tensor<24xf32>, %ub: index) -> tensor<24xf32> {
+ %c0 = arith.constant 0 : index
+ %c6 = arith.constant 6 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %r = scf.for %iv = %c0 to %ub step %c6 iter_args(%acc = %out) -> tensor<24xf32> {
+ %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
+ scf.yield %next : tensor<24xf32>
+ }
+ return %r : tensor<24xf32>
+}
+
+// CHECK-LABEL: func.func @bubble_nonzero_lb(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
+// CHECK: scf.for
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 6] [1, 1]
+func.func @bubble_nonzero_lb(%src: tensor<4x6xf32>, %out: tensor<24xf32>) -> tensor<24xf32> {
+ %c6 = arith.constant 6 : index
+ %c24 = arith.constant 24 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %r = scf.for %iv = %c6 to %c24 step %c6 iter_args(%acc = %out) -> tensor<24xf32> {
+ %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
+ scf.yield %next : tensor<24xf32>
+ }
+ return %r : tensor<24xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_misaligned_lb(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: scf.for
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_misaligned_lb(%src: tensor<4x6xf32>, %out: tensor<24xf32>) -> tensor<24xf32> {
+ %c1 = arith.constant 1 : index
+ %c6 = arith.constant 6 : index
+ %c19 = arith.constant 19 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %r = scf.for %iv = %c1 to %c19 step %c6 iter_args(%acc = %out) -> tensor<24xf32> {
+ %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
+ scf.yield %next : tensor<24xf32>
+ }
+ return %r : tensor<24xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_misaligned_step(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: scf.for
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_misaligned_step(%src: tensor<4x6xf32>, %out: tensor<24xf32>) -> tensor<24xf32> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c19 = arith.constant 19 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %r = scf.for %iv = %c0 to %c19 step %c4 iter_args(%acc = %out) -> tensor<24xf32> {
+ %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
+ scf.yield %next : tensor<24xf32>
+ }
+ return %r : tensor<24xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_dynamic_lb(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: scf.for
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_dynamic_lb(%src: tensor<4x6xf32>, %out: tensor<24xf32>, %lb: index) -> tensor<24xf32> {
+ %c6 = arith.constant 6 : index
+ %c24 = arith.constant 24 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %r = scf.for %iv = %lb to %c24 step %c6 iter_args(%acc = %out) -> tensor<24xf32> {
+ %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
+ scf.yield %next : tensor<24xf32>
+ }
+ return %r : tensor<24xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_dynamic_step(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: scf.for
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_dynamic_step(%src: tensor<4x6xf32>, %out: tensor<24xf32>, %step: index) -> tensor<24xf32> {
+ %c0 = arith.constant 0 : index
+ %c24 = arith.constant 24 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %r = scf.for %iv = %c0 to %c24 step %step iter_args(%acc = %out) -> tensor<24xf32> {
+ %tile = tensor.extract_slice %flat[%iv] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ %next = tensor.insert_slice %tile into %acc[%iv] [6] [1] : tensor<6xf32> into tensor<24xf32>
+ scf.yield %next : tensor<24xf32>
+ }
+ return %r : tensor<24xf32>
+}
+
+// The recurrence happens to be aligned, but this rule only handles IVs.
+// CHECK-LABEL: func.func @no_bubble_iter_arg(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: scf.for
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_iter_arg(%src: tensor<4x6xf32>, %out: tensor<24xf32>) -> tensor<24xf32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c6 = arith.constant 6 : index
+ %flat = tensor.collapse_shape %src [[0, 1]] : tensor<4x6xf32> into tensor<24xf32>
+ %r:2 = scf.for %iv = %c0 to %c4 step %c1 iter_args(%offset = %c0, %acc = %out) -> (index, tensor<24xf32>) {
+ %tile = tensor.extract_slice %flat[%offset] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ %next = tensor.insert_slice %tile into %acc[%offset] [6] [1] : tensor<6xf32> into tensor<24xf32>
+ %next_offset = arith.addi %offset, %c6 : index
+ scf.yield %next_offset, %next : index, tensor<24xf32>
+ }
+ return %r#1 : tensor<24xf32>
+}
+
+// CHECK-LABEL: func.func @no_bubble_arithmetic(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6xf32>
+// CHECK: %[[FLAT:.*]] = tensor.collapse_shape %[[SRC]]
+// CHECK: tensor.extract_slice %[[FLAT]]
+func.func @no_bubble_arithmetic(%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>
+ %tile = tensor.extract_slice %flat[%offset] [6] [1] : tensor<24xf32> to tensor<6xf32>
+ return %tile : tensor<6xf32>
+}
+
+func.func private @consume(tensor<6x6xf32>)
+
+// CHECK-LABEL: func.func @bubble_parallel_ivs(
+// CHECK-SAME: %[[SRC:.*]]: tensor<4x6x4x6xf32>
+// CHECK: scf.parallel
+// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[SRC]]
+// CHECK-SAME: [1, 6, 1, 6] [1, 1, 1, 1]
+func.func @bubble_parallel_ivs(%src: tensor<4x6x4x6xf32>) {
+ %c0 = arith.constant 0 : index
+ %c6 = arith.constant 6 : index
+ %c24 = arith.constant 24 : index
+ %flat = tensor.collapse_shape %src [[0, 1], [2, 3]] : tensor<4x6x4x6xf32> into tensor<24x24xf32>
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c24, %c24) step (%c6, %c6) {
+ %tile = tensor.extract_slice %flat[%i, %j] [6, 6] [1, 1] : tensor<24x24xf32> to tensor<6x6xf32>
+ func.call @consume(%tile) : (tensor<6x6xf32>) -> ()
+ scf.reduce
+ }
+ return
+}
+
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
%func_op = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.op<"func.func">
More information about the Mlir-commits
mailing list