[Mlir-commits] [mlir] [mlir] Fix arith tensor constants (PR #206249)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jul 10 06:27:49 PDT 2026
https://github.com/mygitljf updated https://github.com/llvm/llvm-project/pull/206249
>From 3663396b2155970b6a4982ec42a0108c749df095 Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Sat, 27 Jun 2026 20:57:43 +0000
Subject: [PATCH 1/3] [mlir] Fix arith tensor constants
---
mlir/include/mlir/Conversion/Passes.td | 2 +-
.../mlir/Dialect/Arith/Transforms/Passes.td | 2 +-
.../Conversion/ArithToLLVM/ArithToLLVM.cpp | 3 +-
.../lib/Conversion/ArithToLLVM/CMakeLists.txt | 1 +
.../Dialect/Arith/Transforms/ExpandOps.cpp | 29 +++++++--
.../ArithToLLVM/ceil-floor-div-tensor.mlir | 10 +++
.../test/Conversion/ConvertToSPIRV/arith.mlir | 8 +++
mlir/test/Dialect/Arith/expand-ops.mlir | 63 +++++++++++++++++++
mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp | 2 +
9 files changed, 111 insertions(+), 9 deletions(-)
create mode 100644 mlir/test/Conversion/ArithToLLVM/ceil-floor-div-tensor.mlir
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 550e6b853057d..23463342314b0 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -180,7 +180,7 @@ def ArithToLLVMConversionPass : Pass<"convert-arith-to-llvm"> {
let description = [{
This pass converts supported Arith ops to LLVM dialect instructions.
}];
- let dependentDialects = ["LLVM::LLVMDialect"];
+ let dependentDialects = ["LLVM::LLVMDialect", "tensor::TensorDialect"];
let options = [
Option<"indexBitwidth", "index-bitwidth", "unsigned",
/*default=kDeriveIndexBitwidthFromDataLayout*/"0",
diff --git a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.td b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.td
index 27e9146ec3606..2cee05914d6b5 100644
--- a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.td
@@ -13,7 +13,7 @@ include "mlir/Pass/PassBase.td"
def ArithExpandOpsPass : Pass<"arith-expand"> {
let summary = "Legalize Arith ops to be convertible to LLVM.";
- let dependentDialects = ["vector::VectorDialect"];
+ let dependentDialects = ["tensor::TensorDialect", "vector::VectorDialect"];
let options =
[Option<"includeBf16", "include-bf16", "bool", /*default=*/"false",
"Enable the BF16 expansion patterns">,
diff --git a/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp b/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
index 6bf0fe85bb62f..77e36b45a84af 100644
--- a/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
+++ b/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
@@ -17,6 +17,7 @@
#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/TypeUtilities.h"
#include <type_traits>
@@ -736,7 +737,7 @@ struct ArithToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
: ConvertToLLVMPatternInterface(dialect) {}
void loadDependentDialects(MLIRContext *context) const final {
- context->loadDialect<LLVM::LLVMDialect>();
+ context->loadDialect<LLVM::LLVMDialect, tensor::TensorDialect>();
}
/// Hook for derived dialect interface to provide conversion patterns
diff --git a/mlir/lib/Conversion/ArithToLLVM/CMakeLists.txt b/mlir/lib/Conversion/ArithToLLVM/CMakeLists.txt
index 0a0e25e18b47a..955fcc397dd37 100644
--- a/mlir/lib/Conversion/ArithToLLVM/CMakeLists.txt
+++ b/mlir/lib/Conversion/ArithToLLVM/CMakeLists.txt
@@ -16,4 +16,5 @@ add_mlir_conversion_library(MLIRArithToLLVM
MLIRArithTransforms
MLIRLLVMCommonConversion
MLIRLLVMDialect
+ MLIRTensorDialect
)
diff --git a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
index c9217c57a5f25..d9cfa6bc78dc6 100644
--- a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
@@ -8,6 +8,7 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/Location.h"
@@ -34,6 +35,21 @@ static Value createConst(Location loc, Type type, int value,
return arith::ConstantOp::create(rewriter, loc, attr);
}
+/// Create an integer or index constant, using source as the dynamic shape
+/// source for ranked tensors whose splat cannot be represented as a dense
+/// attribute.
+static Value createConst(Location loc, Type type, int value,
+ PatternRewriter &rewriter, Value source) {
+ auto rankedTensorTy = dyn_cast<RankedTensorType>(type);
+ if (!rankedTensorTy || rankedTensorTy.hasStaticShape())
+ return createConst(loc, type, value, rewriter);
+
+ Value scalar =
+ createConst(loc, rankedTensorTy.getElementType(), value, rewriter);
+ return tensor::SplatOp::create(rewriter, loc, scalar,
+ tensor::getMixedSizes(rewriter, loc, source));
+}
+
/// Create an integer constant from an APInt.
static Value createAPIntConst(Location loc, Type type, const APInt &value,
PatternRewriter &rewriter) {
@@ -76,10 +92,10 @@ struct CeilDivUIOpConverter : public OpRewritePattern<arith::CeilDivUIOp> {
Location loc = op.getLoc();
Value a = op.getLhs();
Value b = op.getRhs();
- Value zero = createConst(loc, a.getType(), 0, rewriter);
+ Value zero = createConst(loc, a.getType(), 0, rewriter, a);
Value compare =
arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq, a, zero);
- Value one = createConst(loc, a.getType(), 1, rewriter);
+ Value one = createConst(loc, a.getType(), 1, rewriter, a);
Value minusOne = arith::SubIOp::create(rewriter, loc, a, one);
Value quotient = arith::DivUIOp::create(rewriter, loc, minusOne, b);
Value plusOne = arith::AddIOp::create(rewriter, loc, quotient, one);
@@ -104,8 +120,8 @@ struct CeilDivSIOpConverter : public OpRewritePattern<arith::CeilDivSIOp> {
Value a = op.getLhs();
Value b = op.getRhs();
- Value zero = createConst(loc, type, 0, rewriter);
- Value one = createConst(loc, type, 1, rewriter);
+ Value zero = createConst(loc, type, 0, rewriter, a);
+ Value one = createConst(loc, type, 1, rewriter, a);
Value quotient = arith::DivSIOp::create(rewriter, loc, a, b);
Value product = arith::MulIOp::create(rewriter, loc, quotient, b);
@@ -150,7 +166,7 @@ struct FloorDivSIOpConverter : public OpRewritePattern<arith::FloorDivSIOp> {
Value product = arith::MulIOp::create(rewriter, loc, quotient, b);
Value notEqualDivisor = arith::CmpIOp::create(
rewriter, loc, arith::CmpIPredicate::ne, a, product);
- Value zero = createConst(loc, type, 0, rewriter);
+ Value zero = createConst(loc, type, 0, rewriter, a);
Value aNeg = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::slt,
a, zero);
@@ -162,7 +178,7 @@ struct FloorDivSIOpConverter : public OpRewritePattern<arith::FloorDivSIOp> {
Value cond =
arith::AndIOp::create(rewriter, loc, notEqualDivisor, signOpposite);
- Value minusOne = createConst(loc, type, -1, rewriter);
+ Value minusOne = createConst(loc, type, -1, rewriter, a);
Value quotientMinusOne =
arith::AddIOp::create(rewriter, loc, quotient, minusOne);
@@ -826,6 +842,7 @@ struct ArithExpandOpsPass
arith::populateArithExpandOpsPatterns(patterns);
target.addLegalDialect<arith::ArithDialect>();
+ target.addLegalDialect<tensor::TensorDialect>();
target.addLegalDialect<vector::VectorDialect>();
// clang-format off
diff --git a/mlir/test/Conversion/ArithToLLVM/ceil-floor-div-tensor.mlir b/mlir/test/Conversion/ArithToLLVM/ceil-floor-div-tensor.mlir
new file mode 100644
index 0000000000000..9d169bc60cade
--- /dev/null
+++ b/mlir/test/Conversion/ArithToLLVM/ceil-floor-div-tensor.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-opt -pass-pipeline="builtin.module(func.func(convert-arith-to-llvm))" %s -split-input-file | FileCheck %s
+// RUN: mlir-opt --convert-to-llvm="filter-dialects=arith" --split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @ceildivui_dynamic_tensor
+// CHECK-SAME: %[[ARG0:.*]]: tensor<8x4x?xi64>) -> tensor<8x4x?xi64>
+func.func @ceildivui_dynamic_tensor(%arg0 : tensor<8x4x?xi64>) -> tensor<8x4x?xi64> {
+// CHECK: arith.ceildivui %[[ARG0]], %[[ARG0]] : tensor<8x4x?xi64>
+ %0 = arith.ceildivui %arg0, %arg0 : tensor<8x4x?xi64>
+ return %0: tensor<8x4x?xi64>
+}
diff --git a/mlir/test/Conversion/ConvertToSPIRV/arith.mlir b/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
index fa204c4df8ace..343c34c9f0ba8 100644
--- a/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
+++ b/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
@@ -89,6 +89,14 @@ func.func @vector_ceildivsi(%lhs: vector<4xi32>, %rhs: vector<4xi32>) -> vector<
return %0 : vector<4xi32>
}
+// CHECK-LABEL: @dynamic_tensor_ceildivui
+// CHECK-SAME: (%[[LHS:.+]]: tensor<8x4x?xi32>, %[[RHS:.+]]: tensor<8x4x?xi32>)
+func.func @dynamic_tensor_ceildivui(%lhs: tensor<8x4x?xi32>, %rhs: tensor<8x4x?xi32>) -> tensor<8x4x?xi32> {
+ // CHECK: arith.ceildivui %[[LHS]], %[[RHS]] : tensor<8x4x?xi32>
+ %0 = arith.ceildivui %lhs, %rhs : tensor<8x4x?xi32>
+ return %0 : tensor<8x4x?xi32>
+}
+
// -----
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Arith/expand-ops.mlir b/mlir/test/Dialect/Arith/expand-ops.mlir
index 75c4de2168761..80b6592417cc2 100644
--- a/mlir/test/Dialect/Arith/expand-ops.mlir
+++ b/mlir/test/Dialect/Arith/expand-ops.mlir
@@ -144,6 +144,69 @@ func.func @ceildivui_index(%arg0: index, %arg1: index) -> (index) {
// -----
+// CHECK-LABEL: func.func @ceildivui_dynamic_tensor(
+// CHECK-SAME: %[[ARG0:.*]]: tensor<8x4x?xi64>,
+// CHECK-SAME: %[[ARG1:.*]]: tensor<8x4x?xi64>) -> tensor<8x4x?xi64> {
+func.func @ceildivui_dynamic_tensor(%arg0: tensor<8x4x?xi64>, %arg1: tensor<8x4x?xi64>) -> tensor<8x4x?xi64> {
+ %res = arith.ceildivui %arg0, %arg1 : tensor<8x4x?xi64>
+ return %res : tensor<8x4x?xi64>
+// CHECK: %[[ZERO_SCALAR:.*]] = arith.constant 0 : i64
+// CHECK: %[[C2_IDX:.*]] = arith.constant 2 : index
+// CHECK: %[[ZERO_DIM:.*]] = tensor.dim %[[ARG0]], %[[C2_IDX]] : tensor<8x4x?xi64>
+// CHECK: %[[ZERO:.*]] = tensor.splat %[[ZERO_SCALAR]]{{\[}}%[[ZERO_DIM]]] : tensor<8x4x?xi64>
+// CHECK: %[[ISZERO:.*]] = arith.cmpi eq, %[[ARG0]], %[[ZERO]] : tensor<8x4x?xi64>
+// CHECK: %[[ONE_SCALAR:.*]] = arith.constant 1 : i64
+// CHECK: %[[C2_IDX_1:.*]] = arith.constant 2 : index
+// CHECK: %[[ONE_DIM:.*]] = tensor.dim %[[ARG0]], %[[C2_IDX_1]] : tensor<8x4x?xi64>
+// CHECK: %[[ONE:.*]] = tensor.splat %[[ONE_SCALAR]]{{\[}}%[[ONE_DIM]]] : tensor<8x4x?xi64>
+// CHECK: %[[SUB:.*]] = arith.subi %[[ARG0]], %[[ONE]] : tensor<8x4x?xi64>
+// CHECK: %[[DIV:.*]] = arith.divui %[[SUB]], %[[ARG1]] : tensor<8x4x?xi64>
+// CHECK: %[[ADD:.*]] = arith.addi %[[DIV]], %[[ONE]] : tensor<8x4x?xi64>
+// CHECK: %[[RES:.*]] = arith.select %[[ISZERO]], %[[ZERO]], %[[ADD]] : tensor<8x4x?xi1>, tensor<8x4x?xi64>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @ceildivsi_dynamic_tensor(
+// CHECK-SAME: %[[ARG0:.*]]: tensor<8x?xi64>,
+// CHECK-SAME: %[[ARG1:.*]]: tensor<8x?xi64>) -> tensor<8x?xi64> {
+func.func @ceildivsi_dynamic_tensor(%arg0: tensor<8x?xi64>, %arg1: tensor<8x?xi64>) -> tensor<8x?xi64> {
+ %res = arith.ceildivsi %arg0, %arg1 : tensor<8x?xi64>
+ return %res : tensor<8x?xi64>
+// CHECK: %[[ZERO_SCALAR:.*]] = arith.constant 0 : i64
+// CHECK: %[[C1_IDX:.*]] = arith.constant 1 : index
+// CHECK: %[[ZERO_DIM:.*]] = tensor.dim %[[ARG0]], %[[C1_IDX]] : tensor<8x?xi64>
+// CHECK: %[[ZERO:.*]] = tensor.splat %[[ZERO_SCALAR]]{{\[}}%[[ZERO_DIM]]] : tensor<8x?xi64>
+// CHECK: %[[ONE_SCALAR:.*]] = arith.constant 1 : i64
+// CHECK: %[[C1_IDX_1:.*]] = arith.constant 1 : index
+// CHECK: %[[ONE_DIM:.*]] = tensor.dim %[[ARG0]], %[[C1_IDX_1]] : tensor<8x?xi64>
+// CHECK: %[[ONE:.*]] = tensor.splat %[[ONE_SCALAR]]{{\[}}%[[ONE_DIM]]] : tensor<8x?xi64>
+// CHECK: %[[DIV:.*]] = arith.divsi %[[ARG0]], %[[ARG1]] : tensor<8x?xi64>
+// CHECK: %[[ADD:.*]] = arith.addi %[[DIV]], %[[ONE]] : tensor<8x?xi64>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @floordivsi_dynamic_tensor(
+// CHECK-SAME: %[[ARG0:.*]]: tensor<?x4xi64>,
+// CHECK-SAME: %[[ARG1:.*]]: tensor<?x4xi64>) -> tensor<?x4xi64> {
+func.func @floordivsi_dynamic_tensor(%arg0: tensor<?x4xi64>, %arg1: tensor<?x4xi64>) -> tensor<?x4xi64> {
+ %res = arith.floordivsi %arg0, %arg1 : tensor<?x4xi64>
+ return %res : tensor<?x4xi64>
+// CHECK: %[[DIV:.*]] = arith.divsi %[[ARG0]], %[[ARG1]] : tensor<?x4xi64>
+// CHECK: %[[ZERO_SCALAR:.*]] = arith.constant 0 : i64
+// CHECK: %[[C0_IDX:.*]] = arith.constant 0 : index
+// CHECK: %[[ZERO_DIM:.*]] = tensor.dim %[[ARG0]], %[[C0_IDX]] : tensor<?x4xi64>
+// CHECK: %[[ZERO:.*]] = tensor.splat %[[ZERO_SCALAR]]{{\[}}%[[ZERO_DIM]]] : tensor<?x4xi64>
+// CHECK: %[[NEG_ONE_SCALAR:.*]] = arith.constant -1 : i64
+// CHECK: %[[C0_IDX_1:.*]] = arith.constant 0 : index
+// CHECK: %[[NEG_ONE_DIM:.*]] = tensor.dim %[[ARG0]], %[[C0_IDX_1]] : tensor<?x4xi64>
+// CHECK: %[[NEG_ONE:.*]] = tensor.splat %[[NEG_ONE_SCALAR]]{{\[}}%[[NEG_ONE_DIM]]] : tensor<?x4xi64>
+// CHECK: %[[SUB_ONE:.*]] = arith.addi %[[DIV]], %[[NEG_ONE]] : tensor<?x4xi64>
+}
+
+// -----
+
// CHECK-LABEL: func @maximumf
func.func @maximumf(%a: f32, %b: f32) -> f32 {
%result = arith.maximumf %a, %b : f32
diff --git a/mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp b/mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp
index 3c99d3c5b60ce..58c66ab6f74c7 100644
--- a/mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp
+++ b/mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp
@@ -18,6 +18,7 @@
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
@@ -93,6 +94,7 @@ struct TestConvertToSPIRVPass final
}
void getDependentDialects(DialectRegistry ®istry) const override {
registry.insert<spirv::SPIRVDialect>();
+ registry.insert<tensor::TensorDialect>();
registry.insert<vector::VectorDialect>();
}
>From 77863de1c79a7c5ae3d6f1b4cc8f62f9caf98498 Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Fri, 10 Jul 2026 20:47:24 +0000
Subject: [PATCH 2/3] [mlir] Refine tensor splat expansion
---
mlir/include/mlir/Conversion/Passes.td | 2 +-
.../mlir/Dialect/Arith/Transforms/Passes.h | 3 +-
.../Conversion/ArithToLLVM/ArithToLLVM.cpp | 3 +-
.../lib/Conversion/ArithToLLVM/CMakeLists.txt | 1 -
.../Dialect/Arith/Transforms/ExpandOps.cpp | 117 ++++++++++++++----
.../ArithToLLVM/ceil-floor-div-tensor.mlir | 10 --
.../test/Conversion/ConvertToSPIRV/arith.mlir | 10 --
mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp | 2 -
8 files changed, 96 insertions(+), 52 deletions(-)
delete mode 100644 mlir/test/Conversion/ArithToLLVM/ceil-floor-div-tensor.mlir
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 23463342314b0..550e6b853057d 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -180,7 +180,7 @@ def ArithToLLVMConversionPass : Pass<"convert-arith-to-llvm"> {
let description = [{
This pass converts supported Arith ops to LLVM dialect instructions.
}];
- let dependentDialects = ["LLVM::LLVMDialect", "tensor::TensorDialect"];
+ let dependentDialects = ["LLVM::LLVMDialect"];
let options = [
Option<"indexBitwidth", "index-bitwidth", "unsigned",
/*default=kDeriveIndexBitwidthFromDataLayout*/"0",
diff --git a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
index 5a07a01d0928a..11081342c2244 100644
--- a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
@@ -54,7 +54,8 @@ void populateEmulateUnsupportedFloatsPatterns(RewritePatternSet &patterns,
void populateEmulateUnsupportedFloatsLegality(ConversionTarget &target,
const TypeConverter &converter);
/// Add patterns to expand Arith ceil/floor division ops.
-void populateCeilFloorDivExpandOpsPatterns(RewritePatternSet &patterns);
+void populateCeilFloorDivExpandOpsPatterns(
+ RewritePatternSet &patterns, bool enableDynamicTensorSplat = false);
/// Add patterns to expand Arith bf16 patterns to lower level bitcasts/shifts.
void populateExpandBFloat16Patterns(RewritePatternSet &patterns);
diff --git a/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp b/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
index 77e36b45a84af..6bf0fe85bb62f 100644
--- a/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
+++ b/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
@@ -17,7 +17,6 @@
#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/TypeUtilities.h"
#include <type_traits>
@@ -737,7 +736,7 @@ struct ArithToLLVMDialectInterface : public ConvertToLLVMPatternInterface {
: ConvertToLLVMPatternInterface(dialect) {}
void loadDependentDialects(MLIRContext *context) const final {
- context->loadDialect<LLVM::LLVMDialect, tensor::TensorDialect>();
+ context->loadDialect<LLVM::LLVMDialect>();
}
/// Hook for derived dialect interface to provide conversion patterns
diff --git a/mlir/lib/Conversion/ArithToLLVM/CMakeLists.txt b/mlir/lib/Conversion/ArithToLLVM/CMakeLists.txt
index 955fcc397dd37..0a0e25e18b47a 100644
--- a/mlir/lib/Conversion/ArithToLLVM/CMakeLists.txt
+++ b/mlir/lib/Conversion/ArithToLLVM/CMakeLists.txt
@@ -16,5 +16,4 @@ add_mlir_conversion_library(MLIRArithToLLVM
MLIRArithTransforms
MLIRLLVMCommonConversion
MLIRLLVMDialect
- MLIRTensorDialect
)
diff --git a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
index d9cfa6bc78dc6..b8c0e39dfddda 100644
--- a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
@@ -24,6 +24,10 @@ namespace arith {
using namespace mlir;
+namespace {
+
+enum class DynamicTensorConstantMaterialization { Disabled, TensorSplat };
+
/// Create an integer or index constant.
static Value createConst(Location loc, Type type, int value,
PatternRewriter &rewriter) {
@@ -38,16 +42,27 @@ static Value createConst(Location loc, Type type, int value,
/// Create an integer or index constant, using source as the dynamic shape
/// source for ranked tensors whose splat cannot be represented as a dense
/// attribute.
-static Value createConst(Location loc, Type type, int value,
- PatternRewriter &rewriter, Value source) {
+static FailureOr<Value>
+createConst(Location loc, Type type, int value, PatternRewriter &rewriter,
+ Value source,
+ DynamicTensorConstantMaterialization dynamicTensorMaterialization) {
auto rankedTensorTy = dyn_cast<RankedTensorType>(type);
+ auto sourceTensorTy = dyn_cast<RankedTensorType>(source.getType());
if (!rankedTensorTy || rankedTensorTy.hasStaticShape())
return createConst(loc, type, value, rewriter);
+ if (!sourceTensorTy ||
+ rankedTensorTy.getShape() != sourceTensorTy.getShape() ||
+ dynamicTensorMaterialization !=
+ DynamicTensorConstantMaterialization::TensorSplat)
+ return failure();
+
+ // Dynamic shaped splats need runtime dimension operands from the source.
Value scalar =
createConst(loc, rankedTensorTy.getElementType(), value, rewriter);
return tensor::SplatOp::create(rewriter, loc, scalar,
- tensor::getMixedSizes(rewriter, loc, source));
+ tensor::getMixedSizes(rewriter, loc, source))
+ .getResult();
}
/// Create an integer constant from an APInt.
@@ -81,27 +96,42 @@ static Type cloneToShapedType(Type cloneFrom, Type cloneTo) {
return cloneTo;
}
-namespace {
-
/// Expands CeilDivUIOp (n, m) into
/// n == 0 ? 0 : ((n-1) / m) + 1
struct CeilDivUIOpConverter : public OpRewritePattern<arith::CeilDivUIOp> {
+ CeilDivUIOpConverter(
+ MLIRContext *context,
+ DynamicTensorConstantMaterialization dynamicTensorMaterialization)
+ : OpRewritePattern<arith::CeilDivUIOp>(context),
+ dynamicTensorMaterialization(dynamicTensorMaterialization) {}
+
using Base::Base;
LogicalResult matchAndRewrite(arith::CeilDivUIOp op,
PatternRewriter &rewriter) const final {
Location loc = op.getLoc();
Value a = op.getLhs();
Value b = op.getRhs();
- Value zero = createConst(loc, a.getType(), 0, rewriter, a);
- Value compare =
- arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq, a, zero);
- Value one = createConst(loc, a.getType(), 1, rewriter, a);
- Value minusOne = arith::SubIOp::create(rewriter, loc, a, one);
+ FailureOr<Value> zero = createConst(loc, a.getType(), 0, rewriter, a,
+ dynamicTensorMaterialization);
+ if (failed(zero))
+ return rewriter.notifyMatchFailure(
+ op, "cannot materialize dynamically shaped tensor constants");
+ Value compare = arith::CmpIOp::create(rewriter, loc,
+ arith::CmpIPredicate::eq, a, *zero);
+ FailureOr<Value> one = createConst(loc, a.getType(), 1, rewriter, a,
+ dynamicTensorMaterialization);
+ if (failed(one))
+ return rewriter.notifyMatchFailure(
+ op, "cannot materialize dynamically shaped tensor constants");
+ Value minusOne = arith::SubIOp::create(rewriter, loc, a, *one);
Value quotient = arith::DivUIOp::create(rewriter, loc, minusOne, b);
- Value plusOne = arith::AddIOp::create(rewriter, loc, quotient, one);
- rewriter.replaceOpWithNewOp<arith::SelectOp>(op, compare, zero, plusOne);
+ Value plusOne = arith::AddIOp::create(rewriter, loc, quotient, *one);
+ rewriter.replaceOpWithNewOp<arith::SelectOp>(op, compare, *zero, plusOne);
return success();
}
+
+private:
+ DynamicTensorConstantMaterialization dynamicTensorMaterialization;
};
/// Expands CeilDivSIOp (a, b) into
@@ -112,6 +142,12 @@ struct CeilDivUIOpConverter : public OpRewritePattern<arith::CeilDivUIOp> {
/// return z;
/// }
struct CeilDivSIOpConverter : public OpRewritePattern<arith::CeilDivSIOp> {
+ CeilDivSIOpConverter(
+ MLIRContext *context,
+ DynamicTensorConstantMaterialization dynamicTensorMaterialization)
+ : OpRewritePattern<arith::CeilDivSIOp>(context),
+ dynamicTensorMaterialization(dynamicTensorMaterialization) {}
+
using Base::Base;
LogicalResult matchAndRewrite(arith::CeilDivSIOp op,
PatternRewriter &rewriter) const final {
@@ -120,8 +156,13 @@ struct CeilDivSIOpConverter : public OpRewritePattern<arith::CeilDivSIOp> {
Value a = op.getLhs();
Value b = op.getRhs();
- Value zero = createConst(loc, type, 0, rewriter, a);
- Value one = createConst(loc, type, 1, rewriter, a);
+ FailureOr<Value> zero =
+ createConst(loc, type, 0, rewriter, a, dynamicTensorMaterialization);
+ FailureOr<Value> one =
+ createConst(loc, type, 1, rewriter, a, dynamicTensorMaterialization);
+ if (failed(zero) || failed(one))
+ return rewriter.notifyMatchFailure(
+ op, "cannot materialize dynamically shaped tensor constants");
Value quotient = arith::DivSIOp::create(rewriter, loc, a, b);
Value product = arith::MulIOp::create(rewriter, loc, quotient, b);
@@ -129,21 +170,25 @@ struct CeilDivSIOpConverter : public OpRewritePattern<arith::CeilDivSIOp> {
rewriter, loc, arith::CmpIPredicate::ne, a, product);
Value aNeg = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::slt,
- a, zero);
+ a, *zero);
Value bNeg = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::slt,
- b, zero);
+ b, *zero);
Value signEqual = arith::CmpIOp::create(
rewriter, loc, arith::CmpIPredicate::eq, aNeg, bNeg);
Value cond =
arith::AndIOp::create(rewriter, loc, notEqualDivisor, signEqual);
- Value quotientPlusOne = arith::AddIOp::create(rewriter, loc, quotient, one);
+ Value quotientPlusOne =
+ arith::AddIOp::create(rewriter, loc, quotient, *one);
rewriter.replaceOpWithNewOp<arith::SelectOp>(op, cond, quotientPlusOne,
quotient);
return success();
}
+
+private:
+ DynamicTensorConstantMaterialization dynamicTensorMaterialization;
};
/// Expands FloorDivSIOp (x, y) into
@@ -154,6 +199,12 @@ struct CeilDivSIOpConverter : public OpRewritePattern<arith::CeilDivSIOp> {
/// return z;
/// }
struct FloorDivSIOpConverter : public OpRewritePattern<arith::FloorDivSIOp> {
+ FloorDivSIOpConverter(
+ MLIRContext *context,
+ DynamicTensorConstantMaterialization dynamicTensorMaterialization)
+ : OpRewritePattern<arith::FloorDivSIOp>(context),
+ dynamicTensorMaterialization(dynamicTensorMaterialization) {}
+
using Base::Base;
LogicalResult matchAndRewrite(arith::FloorDivSIOp op,
PatternRewriter &rewriter) const final {
@@ -166,26 +217,37 @@ struct FloorDivSIOpConverter : public OpRewritePattern<arith::FloorDivSIOp> {
Value product = arith::MulIOp::create(rewriter, loc, quotient, b);
Value notEqualDivisor = arith::CmpIOp::create(
rewriter, loc, arith::CmpIPredicate::ne, a, product);
- Value zero = createConst(loc, type, 0, rewriter, a);
+ FailureOr<Value> zero =
+ createConst(loc, type, 0, rewriter, a, dynamicTensorMaterialization);
+ if (failed(zero))
+ return rewriter.notifyMatchFailure(
+ op, "cannot materialize dynamically shaped tensor constants");
Value aNeg = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::slt,
- a, zero);
+ a, *zero);
Value bNeg = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::slt,
- b, zero);
+ b, *zero);
Value signOpposite = arith::CmpIOp::create(
rewriter, loc, arith::CmpIPredicate::ne, aNeg, bNeg);
Value cond =
arith::AndIOp::create(rewriter, loc, notEqualDivisor, signOpposite);
- Value minusOne = createConst(loc, type, -1, rewriter, a);
+ FailureOr<Value> minusOne =
+ createConst(loc, type, -1, rewriter, a, dynamicTensorMaterialization);
+ if (failed(minusOne))
+ return rewriter.notifyMatchFailure(
+ op, "cannot materialize dynamically shaped tensor constants");
Value quotientMinusOne =
- arith::AddIOp::create(rewriter, loc, quotient, minusOne);
+ arith::AddIOp::create(rewriter, loc, quotient, *minusOne);
rewriter.replaceOpWithNewOp<arith::SelectOp>(op, cond, quotientMinusOne,
quotient);
return success();
}
+
+private:
+ DynamicTensorConstantMaterialization dynamicTensorMaterialization;
};
template <typename OpTy, arith::CmpIPredicate pred>
@@ -921,10 +983,14 @@ struct ArithExpandOpsPass
} // namespace
void mlir::arith::populateCeilFloorDivExpandOpsPatterns(
- RewritePatternSet &patterns) {
+ RewritePatternSet &patterns, bool enableDynamicTensorSplat) {
+ DynamicTensorConstantMaterialization dynamicTensorMaterialization =
+ enableDynamicTensorSplat
+ ? DynamicTensorConstantMaterialization::TensorSplat
+ : DynamicTensorConstantMaterialization::Disabled;
patterns
.add<CeilDivSIOpConverter, CeilDivUIOpConverter, FloorDivSIOpConverter>(
- patterns.getContext());
+ patterns.getContext(), dynamicTensorMaterialization);
}
void mlir::arith::populateExpandBFloat16Patterns(RewritePatternSet &patterns) {
@@ -954,7 +1020,8 @@ void mlir::arith::populateExpandFlushDenormalsPatterns(
}
void mlir::arith::populateArithExpandOpsPatterns(RewritePatternSet &patterns) {
- populateCeilFloorDivExpandOpsPatterns(patterns);
+ populateCeilFloorDivExpandOpsPatterns(patterns,
+ /*enableDynamicTensorSplat=*/true);
populateExpandScalingExtTruncPatterns(patterns);
// clang-format off
patterns.add<
diff --git a/mlir/test/Conversion/ArithToLLVM/ceil-floor-div-tensor.mlir b/mlir/test/Conversion/ArithToLLVM/ceil-floor-div-tensor.mlir
deleted file mode 100644
index 9d169bc60cade..0000000000000
--- a/mlir/test/Conversion/ArithToLLVM/ceil-floor-div-tensor.mlir
+++ /dev/null
@@ -1,10 +0,0 @@
-// RUN: mlir-opt -pass-pipeline="builtin.module(func.func(convert-arith-to-llvm))" %s -split-input-file | FileCheck %s
-// RUN: mlir-opt --convert-to-llvm="filter-dialects=arith" --split-input-file %s | FileCheck %s
-
-// CHECK-LABEL: @ceildivui_dynamic_tensor
-// CHECK-SAME: %[[ARG0:.*]]: tensor<8x4x?xi64>) -> tensor<8x4x?xi64>
-func.func @ceildivui_dynamic_tensor(%arg0 : tensor<8x4x?xi64>) -> tensor<8x4x?xi64> {
-// CHECK: arith.ceildivui %[[ARG0]], %[[ARG0]] : tensor<8x4x?xi64>
- %0 = arith.ceildivui %arg0, %arg0 : tensor<8x4x?xi64>
- return %0: tensor<8x4x?xi64>
-}
diff --git a/mlir/test/Conversion/ConvertToSPIRV/arith.mlir b/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
index 343c34c9f0ba8..65a8180e480b7 100644
--- a/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
+++ b/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
@@ -89,16 +89,6 @@ func.func @vector_ceildivsi(%lhs: vector<4xi32>, %rhs: vector<4xi32>) -> vector<
return %0 : vector<4xi32>
}
-// CHECK-LABEL: @dynamic_tensor_ceildivui
-// CHECK-SAME: (%[[LHS:.+]]: tensor<8x4x?xi32>, %[[RHS:.+]]: tensor<8x4x?xi32>)
-func.func @dynamic_tensor_ceildivui(%lhs: tensor<8x4x?xi32>, %rhs: tensor<8x4x?xi32>) -> tensor<8x4x?xi32> {
- // CHECK: arith.ceildivui %[[LHS]], %[[RHS]] : tensor<8x4x?xi32>
- %0 = arith.ceildivui %lhs, %rhs : tensor<8x4x?xi32>
- return %0 : tensor<8x4x?xi32>
-}
-
-// -----
-
//===----------------------------------------------------------------------===//
// arith bit ops
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp b/mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp
index 58c66ab6f74c7..3c99d3c5b60ce 100644
--- a/mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp
+++ b/mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp
@@ -18,7 +18,6 @@
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
@@ -94,7 +93,6 @@ struct TestConvertToSPIRVPass final
}
void getDependentDialects(DialectRegistry ®istry) const override {
registry.insert<spirv::SPIRVDialect>();
- registry.insert<tensor::TensorDialect>();
registry.insert<vector::VectorDialect>();
}
>From ae1131363388ffca29bb2d1c70e2d379a716701c Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Fri, 10 Jul 2026 21:30:10 +0000
Subject: [PATCH 3/3] [mlir] Restore SPIR-V test
---
mlir/test/Conversion/ConvertToSPIRV/arith.mlir | 10 ++++++++++
mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp | 2 ++
2 files changed, 12 insertions(+)
diff --git a/mlir/test/Conversion/ConvertToSPIRV/arith.mlir b/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
index 65a8180e480b7..343c34c9f0ba8 100644
--- a/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
+++ b/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
@@ -89,6 +89,16 @@ func.func @vector_ceildivsi(%lhs: vector<4xi32>, %rhs: vector<4xi32>) -> vector<
return %0 : vector<4xi32>
}
+// CHECK-LABEL: @dynamic_tensor_ceildivui
+// CHECK-SAME: (%[[LHS:.+]]: tensor<8x4x?xi32>, %[[RHS:.+]]: tensor<8x4x?xi32>)
+func.func @dynamic_tensor_ceildivui(%lhs: tensor<8x4x?xi32>, %rhs: tensor<8x4x?xi32>) -> tensor<8x4x?xi32> {
+ // CHECK: arith.ceildivui %[[LHS]], %[[RHS]] : tensor<8x4x?xi32>
+ %0 = arith.ceildivui %lhs, %rhs : tensor<8x4x?xi32>
+ return %0 : tensor<8x4x?xi32>
+}
+
+// -----
+
//===----------------------------------------------------------------------===//
// arith bit ops
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp b/mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp
index 3c99d3c5b60ce..58c66ab6f74c7 100644
--- a/mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp
+++ b/mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp
@@ -18,6 +18,7 @@
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
@@ -93,6 +94,7 @@ struct TestConvertToSPIRVPass final
}
void getDependentDialects(DialectRegistry ®istry) const override {
registry.insert<spirv::SPIRVDialect>();
+ registry.insert<tensor::TensorDialect>();
registry.insert<vector::VectorDialect>();
}
More information about the Mlir-commits
mailing list