[Mlir-commits] [mlir] [mlir] Fix arith tensor constants (PR #206249)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Jun 27 06:34:05 PDT 2026
https://github.com/mygitljf created https://github.com/llvm/llvm-project/pull/206249
I updated the arith ceil/floor division expansion so dynamic ranked tensor constants are materialized through scalar constants plus tensor.splat using the source tensor shape. Static tensors, vectors, and scalar cases still use the existing constant path, so this keeps the old behavior where dense constants are valid.
I also added the needed Tensor dialect dependencies for the pass and conversion entry points that can now build tensor.dim/tensor.splat while expanding these ops. The regression coverage includes arith-expand, direct ArithToLLVM conversion, and the SPIR-V test conversion path that reuses the same expansion patterns.
Fixes #206091
>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] [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>();
}
More information about the Mlir-commits
mailing list