[Mlir-commits] [mlir] [mlir] Fix arith tensor constants (PR #206249)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jul 14 01:52:23 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/8] [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/8] [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/8] [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>();
}
>From 4f6a92bef6a65f92a103adb2c4a6c8211ea31352 Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Fri, 10 Jul 2026 21:54:16 +0000
Subject: [PATCH 4/8] [mlir] Drop SPIR-V tensor test
---
mlir/test/Conversion/ConvertToSPIRV/arith.mlir | 8 --------
mlir/test/lib/Pass/TestConvertToSPIRVPass.cpp | 2 --
2 files changed, 10 deletions(-)
diff --git a/mlir/test/Conversion/ConvertToSPIRV/arith.mlir b/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
index 343c34c9f0ba8..fa204c4df8ace 100644
--- a/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
+++ b/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
@@ -89,14 +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>
-}
-
// -----
//===----------------------------------------------------------------------===//
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 2d6c5dced435f347fb4b604e2a4b12c73a395502 Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Tue, 14 Jul 2026 16:49:52 +0800
Subject: [PATCH 5/8] [MLIR][Arith] Generalize division expansion
---
.../mlir/Dialect/Arith/Transforms/Passes.h | 3 +-
.../mlir/Dialect/Arith/Transforms/Passes.td | 2 +-
.../Dialect/Arith/Transforms/ExpandOps.cpp | 229 ++-----
mlir/test/Dialect/Arith/expand-ops.mlir | 561 +++++++++++++-----
4 files changed, 464 insertions(+), 331 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
index 11081342c2244..5a07a01d0928a 100644
--- a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
@@ -54,8 +54,7 @@ void populateEmulateUnsupportedFloatsPatterns(RewritePatternSet &patterns,
void populateEmulateUnsupportedFloatsLegality(ConversionTarget &target,
const TypeConverter &converter);
/// Add patterns to expand Arith ceil/floor division ops.
-void populateCeilFloorDivExpandOpsPatterns(
- RewritePatternSet &patterns, bool enableDynamicTensorSplat = false);
+void populateCeilFloorDivExpandOpsPatterns(RewritePatternSet &patterns);
/// Add patterns to expand Arith bf16 patterns to lower level bitcasts/shifts.
void populateExpandBFloat16Patterns(RewritePatternSet &patterns);
diff --git a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.td b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.td
index 2cee05914d6b5..27e9146ec3606 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 = ["tensor::TensorDialect", "vector::VectorDialect"];
+ let dependentDialects = ["vector::VectorDialect"];
let options =
[Option<"includeBf16", "include-bf16", "bool", /*default=*/"false",
"Enable the BF16 expansion patterns">,
diff --git a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
index b8c0e39dfddda..e165bb2f2c0f1 100644
--- a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
@@ -8,7 +8,6 @@
#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"
@@ -24,10 +23,6 @@ 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) {
@@ -39,32 +34,6 @@ 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 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))
- .getResult();
-}
-
/// Create an integer constant from an APInt.
static Value createAPIntConst(Location loc, Type type, const APInt &value,
PatternRewriter &rewriter) {
@@ -96,158 +65,94 @@ static Type cloneToShapedType(Type cloneFrom, Type cloneTo) {
return cloneTo;
}
-/// 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) {}
+static Value extendBoolToType(Location loc, Value value, Type type,
+ PatternRewriter &rewriter) {
+ if (value.getType() == type)
+ return value;
+ if (getElementTypeOrSelf(type).isIndex())
+ return arith::IndexCastUIOp::create(rewriter, loc, type, value);
+ return arith::ExtUIOp::create(rewriter, loc, type, value);
+}
+
+namespace {
+/// Expands CeilDivUIOp (lhs, rhs) into
+/// q = lhs / rhs
+/// q * rhs != lhs ? q + 1 : q
+struct CeilDivUIOpConverter : public OpRewritePattern<arith::CeilDivUIOp> {
using Base::Base;
LogicalResult matchAndRewrite(arith::CeilDivUIOp op,
PatternRewriter &rewriter) const final {
Location loc = op.getLoc();
- Value a = op.getLhs();
- Value b = op.getRhs();
- 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 lhs = op.getLhs();
+ Value rhs = op.getRhs();
+
+ Value quotient = arith::DivUIOp::create(rewriter, loc, lhs, rhs);
+ Value product = arith::MulIOp::create(rewriter, loc, quotient, rhs);
+ Value inexact = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::ne, product, lhs);
+ Value adjustment = extendBoolToType(loc, inexact, op.getType(), rewriter);
+ rewriter.replaceOpWithNewOp<arith::AddIOp>(op, quotient, adjustment);
return success();
}
-
-private:
- DynamicTensorConstantMaterialization dynamicTensorMaterialization;
};
-/// Expands CeilDivSIOp (a, b) into
-/// z = a / b
-/// if (z * b != a && (a < 0) == (b < 0)) {
-/// return z + 1;
-/// } else {
-/// return z;
-/// }
+/// Expands CeilDivSIOp (lhs, rhs) by rounding an inexact quotient when the
+/// signed and unsigned operand orderings agree (the operands have equal sign).
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 {
Location loc = op.getLoc();
- Type type = op.getType();
- Value a = op.getLhs();
- Value b = op.getRhs();
-
- 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);
- Value notEqualDivisor = arith::CmpIOp::create(
- rewriter, loc, arith::CmpIPredicate::ne, a, product);
-
- Value aNeg = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::slt,
- a, *zero);
- Value bNeg = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::slt,
- 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 lhs = op.getLhs();
+ Value rhs = op.getRhs();
- rewriter.replaceOpWithNewOp<arith::SelectOp>(op, cond, quotientPlusOne,
- quotient);
+ Value quotient = arith::DivSIOp::create(rewriter, loc, lhs, rhs);
+ Value product = arith::MulIOp::create(rewriter, loc, quotient, rhs);
+ Value inexact = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::ne, product, lhs);
+ Value signedLess = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::slt, lhs, rhs);
+ Value unsignedLess = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::ult, lhs, rhs);
+ Value sameSign = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::eq, signedLess, unsignedLess);
+ Value shouldRound = arith::AndIOp::create(rewriter, loc, inexact, sameSign);
+ Value adjustment =
+ extendBoolToType(loc, shouldRound, op.getType(), rewriter);
+ rewriter.replaceOpWithNewOp<arith::AddIOp>(op, quotient, adjustment);
return success();
}
-
-private:
- DynamicTensorConstantMaterialization dynamicTensorMaterialization;
};
-/// Expands FloorDivSIOp (x, y) into
-/// z = x / y
-/// if (z * y != x && (x < 0) != (y < 0)) {
-/// return z - 1;
-/// } else {
-/// return z;
-/// }
+/// Expands FloorDivSIOp (lhs, rhs) by rounding an inexact quotient when the
+/// signed and unsigned operand orderings differ (the operands have opposite
+/// signs).
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 {
Location loc = op.getLoc();
- Type type = op.getType();
- Value a = op.getLhs();
- Value b = op.getRhs();
-
- Value quotient = arith::DivSIOp::create(rewriter, loc, a, b);
- Value product = arith::MulIOp::create(rewriter, loc, quotient, b);
- Value notEqualDivisor = arith::CmpIOp::create(
- rewriter, loc, arith::CmpIPredicate::ne, a, product);
- 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);
- Value bNeg = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::slt,
- b, *zero);
-
- Value signOpposite = arith::CmpIOp::create(
- rewriter, loc, arith::CmpIPredicate::ne, aNeg, bNeg);
- Value cond =
- arith::AndIOp::create(rewriter, loc, notEqualDivisor, signOpposite);
-
- 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);
+ Value lhs = op.getLhs();
+ Value rhs = op.getRhs();
- rewriter.replaceOpWithNewOp<arith::SelectOp>(op, cond, quotientMinusOne,
- quotient);
+ Value quotient = arith::DivSIOp::create(rewriter, loc, lhs, rhs);
+ Value product = arith::MulIOp::create(rewriter, loc, quotient, rhs);
+ Value inexact = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::ne, product, lhs);
+ Value signedLess = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::slt, lhs, rhs);
+ Value unsignedLess = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::ult, lhs, rhs);
+ Value oppositeSign = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::ne, signedLess, unsignedLess);
+ Value shouldRound =
+ arith::AndIOp::create(rewriter, loc, inexact, oppositeSign);
+ Value adjustment =
+ extendBoolToType(loc, shouldRound, op.getType(), rewriter);
+ rewriter.replaceOpWithNewOp<arith::SubIOp>(op, quotient, adjustment);
return success();
}
-
-private:
- DynamicTensorConstantMaterialization dynamicTensorMaterialization;
};
template <typename OpTy, arith::CmpIPredicate pred>
@@ -904,7 +809,6 @@ struct ArithExpandOpsPass
arith::populateArithExpandOpsPatterns(patterns);
target.addLegalDialect<arith::ArithDialect>();
- target.addLegalDialect<tensor::TensorDialect>();
target.addLegalDialect<vector::VectorDialect>();
// clang-format off
@@ -983,14 +887,10 @@ struct ArithExpandOpsPass
} // namespace
void mlir::arith::populateCeilFloorDivExpandOpsPatterns(
- RewritePatternSet &patterns, bool enableDynamicTensorSplat) {
- DynamicTensorConstantMaterialization dynamicTensorMaterialization =
- enableDynamicTensorSplat
- ? DynamicTensorConstantMaterialization::TensorSplat
- : DynamicTensorConstantMaterialization::Disabled;
+ RewritePatternSet &patterns) {
patterns
.add<CeilDivSIOpConverter, CeilDivUIOpConverter, FloorDivSIOpConverter>(
- patterns.getContext(), dynamicTensorMaterialization);
+ patterns.getContext());
}
void mlir::arith::populateExpandBFloat16Patterns(RewritePatternSet &patterns) {
@@ -1020,8 +920,7 @@ void mlir::arith::populateExpandFlushDenormalsPatterns(
}
void mlir::arith::populateArithExpandOpsPatterns(RewritePatternSet &patterns) {
- populateCeilFloorDivExpandOpsPatterns(patterns,
- /*enableDynamicTensorSplat=*/true);
+ populateCeilFloorDivExpandOpsPatterns(patterns);
populateExpandScalingExtTruncPatterns(patterns);
// clang-format off
patterns.add<
diff --git a/mlir/test/Dialect/Arith/expand-ops.mlir b/mlir/test/Dialect/Arith/expand-ops.mlir
index 80b6592417cc2..f8432a6a78627 100644
--- a/mlir/test/Dialect/Arith/expand-ops.mlir
+++ b/mlir/test/Dialect/Arith/expand-ops.mlir
@@ -1,208 +1,443 @@
// RUN: mlir-opt %s -arith-expand="include-bf16=true include-f8e8m0=true include-f4e2m1=true" -verify-diagnostics -split-input-file | FileCheck %s
// RUN: mlir-opt %s -arith-expand -split-input-file -verify-diagnostics | FileCheck %s --check-prefix=SCHECK
-
-// Test ceil divide with signed integer
-// CHECK-LABEL: func @ceildivi
-// CHECK-SAME: ([[ARG0:%.+]]: i32, [[ARG1:%.+]]: i32) -> i32 {
-func.func @ceildivi(%arg0: i32, %arg1: i32) -> (i32) {
- %res = arith.ceildivsi %arg0, %arg1 : i32
+// RUN: mlir-opt %s -arith-expand="include-bf16=true include-f8e8m0=true include-f4e2m1=true" -canonicalize -verify-diagnostics -split-input-file | FileCheck %s --check-prefix=VALUES
+
+// CHECK-LABEL: func.func @ceildivui(
+// CHECK-SAME: %[[LHS:.*]]: i32, %[[RHS:.*]]: i32) -> i32 {
+// CHECK-NOT: arith.ceildivui
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divui %[[LHS]], %[[RHS]] : i32
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : i32
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : i32
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.extui %[[INEXACT]] : i1 to i32
+// CHECK-NEXT: %[[RESULT:.*]] = arith.addi %[[Q]], %[[ADJUSTMENT]] : i32
+// CHECK-NEXT: return %[[RESULT]] : i32
+func.func @ceildivui(%arg0: i32, %arg1: i32) -> i32 {
+ %res = arith.ceildivui %arg0, %arg1 : i32
return %res : i32
-
-// CHECK: [[ZERO:%.+]] = arith.constant 0 : i32
-// CHECK: [[ONE:%.+]] = arith.constant 1 : i32
-// CHECK: [[DIV:%.+]] = arith.divsi %arg0, %arg1 : i32
-// CHECK: [[MUL:%.+]] = arith.muli [[DIV]], %arg1 : i32
-// CHECK: [[NEXACT:%.+]] = arith.cmpi ne, %arg0, [[MUL]] : i32
-// CHECK: [[NNEG:%.+]] = arith.cmpi slt, %arg0, [[ZERO]] : i32
-// CHECK: [[MNEG:%.+]] = arith.cmpi slt, %arg1, [[ZERO]] : i32
-// CHECK: [[SAMESIGN:%.+]] = arith.cmpi eq, [[NNEG]], [[MNEG]] : i1
-// CHECK: [[SHOULDROUND:%.+]] = arith.andi [[NEXACT]], [[SAMESIGN]] : i1
-// CHECK: [[CEIL:%.+]] = arith.addi [[DIV]], [[ONE]] : i32
-// CHECK: [[RES:%.+]] = arith.select [[SHOULDROUND]], [[CEIL]], [[DIV]] : i32
}
// -----
-// Test ceil divide with index type
-// CHECK-LABEL: func @ceildivi_index
-// CHECK-SAME: ([[ARG0:%.+]]: index, [[ARG1:%.+]]: index) -> index {
-func.func @ceildivi_index(%arg0: index, %arg1: index) -> (index) {
- %res = arith.ceildivsi %arg0, %arg1 : index
+// CHECK-LABEL: func.func @ceildivui_index(
+// CHECK-SAME: %[[LHS:.*]]: index, %[[RHS:.*]]: index) -> index {
+// CHECK-NOT: arith.ceildivui
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.extui
+// CHECK-NEXT: %[[Q:.*]] = arith.divui %[[LHS]], %[[RHS]] : index
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : index
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : index
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.index_castui %[[INEXACT]] : i1 to index
+// CHECK-NEXT: %[[RESULT:.*]] = arith.addi %[[Q]], %[[ADJUSTMENT]] : index
+// CHECK-NEXT: return %[[RESULT]] : index
+func.func @ceildivui_index(%arg0: index, %arg1: index) -> index {
+ %res = arith.ceildivui %arg0, %arg1 : index
return %res : index
+}
-// CHECK: [[ZERO:%.+]] = arith.constant 0 : index
-// CHECK: [[ONE:%.+]] = arith.constant 1 : index
-// CHECK: [[DIV:%.+]] = arith.divsi %arg0, %arg1 : index
-// CHECK: [[MUL:%.+]] = arith.muli [[DIV]], %arg1 : index
-// CHECK: [[NEXACT:%.+]] = arith.cmpi ne, %arg0, [[MUL]] : index
-// CHECK: [[NNEG:%.+]] = arith.cmpi slt, %arg0, [[ZERO]] : index
-// CHECK: [[MNEG:%.+]] = arith.cmpi slt, %arg1, [[ZERO]] : index
-// CHECK: [[SAMESIGN:%.+]] = arith.cmpi eq, [[NNEG]], [[MNEG]] : i1
-// CHECK: [[SHOULDROUND:%.+]] = arith.andi [[NEXACT]], [[SAMESIGN]] : i1
-// CHECK: [[CEIL:%.+]] = arith.addi [[DIV]], [[ONE]] : index
-// CHECK: [[RES:%.+]] = arith.select [[SHOULDROUND]], [[CEIL]], [[DIV]] : index
+// -----
+// CHECK-LABEL: func.func @ceildivui_vec(
+// CHECK-SAME: %[[LHS:.*]]: vector<4xi32>, %[[RHS:.*]]: vector<4xi32>) -> vector<4xi32> {
+// CHECK-NOT: arith.ceildivui
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divui %[[LHS]], %[[RHS]] : vector<4xi32>
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : vector<4xi32>
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : vector<4xi32>
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.extui %[[INEXACT]] : vector<4xi1> to vector<4xi32>
+// CHECK-NEXT: %[[RESULT:.*]] = arith.addi %[[Q]], %[[ADJUSTMENT]] : vector<4xi32>
+// CHECK-NEXT: return %[[RESULT]] : vector<4xi32>
+func.func @ceildivui_vec(%arg0: vector<4xi32>, %arg1: vector<4xi32>) -> vector<4xi32> {
+ %res = arith.ceildivui %arg0, %arg1 : vector<4xi32>
+ return %res : vector<4xi32>
}
// -----
-// Test floor divide with signed integer
-// CHECK-LABEL: func @floordivi
-// CHECK-SAME: ([[ARG0:%.+]]: i32, [[ARG1:%.+]]: i32) -> i32 {
-func.func @floordivi(%arg0: i32, %arg1: i32) -> (i32) {
- %res = arith.floordivsi %arg0, %arg1 : i32
- return %res : i32
-// CHECK: %[[QUOTIENT:.*]] = arith.divsi %arg0, %arg1 : i32
-// CHECK: %[[PRODUCT:.*]] = arith.muli %[[QUOTIENT]], %arg1 : i32
-// CHECK: %[[NOT_EQ_PRODUCT:.*]] = arith.cmpi ne, %arg0, %[[PRODUCT]] : i32
-// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0 : i32
-// CHECK: %[[NEG_DIVISOR:.*]] = arith.cmpi slt, %arg0, %[[ZERO]] : i32
-// CHECK: %[[NEG_DIVIDEND:.*]] = arith.cmpi slt, %arg1, %[[ZERO]] : i32
-// CHECK: %[[OPPOSITE_SIGN:.*]] = arith.cmpi ne, %[[NEG_DIVISOR]], %[[NEG_DIVIDEND]] : i1
-// CHECK: %[[CONDITION:.*]] = arith.andi %[[NOT_EQ_PRODUCT]], %[[OPPOSITE_SIGN]] : i1
-// CHECK-DAG: %[[NEG_ONE:.*]] = arith.constant -1 : i32
-// CHECK: %[[MINUS_ONE:.*]] = arith.addi %[[QUOTIENT]], %[[NEG_ONE]] : i32
-// CHECK: %[[RES:.*]] = arith.select %[[CONDITION]], %[[MINUS_ONE]], %[[QUOTIENT]] : i32
+// CHECK-LABEL: func.func @ceildivui_static_tensor(
+// CHECK-SAME: %[[LHS:.*]]: tensor<2x3xi32>, %[[RHS:.*]]: tensor<2x3xi32>) -> tensor<2x3xi32> {
+// CHECK-NOT: arith.ceildivui
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divui %[[LHS]], %[[RHS]] : tensor<2x3xi32>
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : tensor<2x3xi32>
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : tensor<2x3xi32>
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.extui %[[INEXACT]] : tensor<2x3xi1> to tensor<2x3xi32>
+// CHECK-NEXT: %[[RESULT:.*]] = arith.addi %[[Q]], %[[ADJUSTMENT]] : tensor<2x3xi32>
+// CHECK-NEXT: return %[[RESULT]] : tensor<2x3xi32>
+func.func @ceildivui_static_tensor(%arg0: tensor<2x3xi32>, %arg1: tensor<2x3xi32>) -> tensor<2x3xi32> {
+ %res = arith.ceildivui %arg0, %arg1 : tensor<2x3xi32>
+ return %res : tensor<2x3xi32>
}
// -----
-// Test floor divide with index type
-// CHECK-LABEL: func @floordivi_index
-// CHECK-SAME: ([[ARG0:%.+]]: index, [[ARG1:%.+]]: index) -> index {
-func.func @floordivi_index(%arg0: index, %arg1: index) -> (index) {
- %res = arith.floordivsi %arg0, %arg1 : index
- return %res : index
-// CHECK: %[[QUOTIENT:.*]] = arith.divsi %arg0, %arg1 : index
-// CHECK: %[[PRODUCT:.*]] = arith.muli %[[QUOTIENT]], %arg1 : index
-// CHECK: %[[NOT_EQ_PRODUCT:.*]] = arith.cmpi ne, %arg0, %[[PRODUCT]] : index
-// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0 : index
-// CHECK: %[[NEG_DIVISOR:.*]] = arith.cmpi slt, %arg0, %[[ZERO]] : index
-// CHECK: %[[NEG_DIVIDEND:.*]] = arith.cmpi slt, %arg1, %[[ZERO]] : index
-// CHECK: %[[OPPOSITE_SIGN:.*]] = arith.cmpi ne, %[[NEG_DIVISOR]], %[[NEG_DIVIDEND]] : i1
-// CHECK: %[[CONDITION:.*]] = arith.andi %[[NOT_EQ_PRODUCT]], %[[OPPOSITE_SIGN]] : i1
-// CHECK: %[[NEG_ONE:.*]] = arith.constant -1 : index
-// CHECK-DAG: %[[MINUS_ONE:.*]] = arith.addi %[[QUOTIENT]], %[[NEG_ONE]] : index
-// CHECK: %[[RES:.*]] = arith.select %[[CONDITION]], %[[MINUS_ONE]], %[[QUOTIENT]] : index
-}
-
-// -----
-
-// Test floor divide with vector
-// CHECK-LABEL: func.func @floordivi_vec(
-// CHECK-SAME: %[[VAL_0:.*]]: vector<4xi32>,
-// CHECK-SAME: %[[VAL_1:.*]]: vector<4xi32>) -> vector<4xi32> {
-func.func @floordivi_vec(%arg0: vector<4xi32>, %arg1: vector<4xi32>) -> (vector<4xi32>) {
- %res = arith.floordivsi %arg0, %arg1 : vector<4xi32>
- return %res : vector<4xi32>
-// CHECK: %[[QUOTIENT:.*]] = arith.divsi %arg0, %arg1 : vector<4xi32>
-// CHECK: %[[PRODUCT:.*]] = arith.muli %[[QUOTIENT]], %arg1 : vector<4xi32>
-// CHECK: %[[NOT_EQ_PRODUCT:.*]] = arith.cmpi ne, %arg0, %[[PRODUCT]] : vector<4xi32>
-// CHECK-DAG: %[[ZERO:.*]] = arith.constant dense<0> : vector<4xi32>
-// CHECK: %[[NEG_DIVISOR:.*]] = arith.cmpi slt, %arg0, %[[ZERO]] : vector<4xi32>
-// CHECK: %[[NEG_DIVIDEND:.*]] = arith.cmpi slt, %arg1, %[[ZERO]] : vector<4xi32>
-// CHECK: %[[OPPOSITE_SIGN:.*]] = arith.cmpi ne, %[[NEG_DIVISOR]], %[[NEG_DIVIDEND]] : vector<4xi1>
-// CHECK: %[[CONDITION:.*]] = arith.andi %[[NOT_EQ_PRODUCT]], %[[OPPOSITE_SIGN]] : vector<4xi1>
-// CHECK-DAG: %[[NEG_ONE:.*]] = arith.constant dense<-1> : vector<4xi32>
-// CHECK: %[[MINUS_ONE:.*]] = arith.addi %[[QUOTIENT]], %[[NEG_ONE]] : vector<4xi32>
-// CHECK: %[[RES:.*]] = arith.select %[[CONDITION]], %[[MINUS_ONE]], %[[QUOTIENT]] : vector<4xi1>, vector<4xi32>
+// CHECK-LABEL: func.func @ceildivui_dynamic_tensor(
+// CHECK-SAME: %[[LHS:.*]]: tensor<8x4x?xi64>, %[[RHS:.*]]: tensor<8x4x?xi64>) -> tensor<8x4x?xi64> {
+// CHECK-NOT: arith.ceildivui
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divui %[[LHS]], %[[RHS]] : tensor<8x4x?xi64>
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : tensor<8x4x?xi64>
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : tensor<8x4x?xi64>
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.extui %[[INEXACT]] : tensor<8x4x?xi1> to tensor<8x4x?xi64>
+// CHECK-NEXT: %[[RESULT:.*]] = arith.addi %[[Q]], %[[ADJUSTMENT]] : tensor<8x4x?xi64>
+// CHECK-NEXT: return %[[RESULT]] : 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>
}
// -----
-// Test ceil divide with unsigned integer
-// CHECK-LABEL: func @ceildivui
-// CHECK-SAME: ([[ARG0:%.+]]: i32, [[ARG1:%.+]]: i32) -> i32 {
-func.func @ceildivui(%arg0: i32, %arg1: i32) -> (i32) {
- %res = arith.ceildivui %arg0, %arg1 : i32
+// CHECK-LABEL: func.func @ceildivui_i1(
+// CHECK-SAME: %[[LHS:.*]]: i1, %[[RHS:.*]]: i1) -> i1 {
+// CHECK-NOT: arith.ceildivui
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.extui
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divui %[[LHS]], %[[RHS]] : i1
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : i1
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : i1
+// CHECK-NEXT: %[[RESULT:.*]] = arith.addi %[[Q]], %[[INEXACT]] : i1
+// CHECK-NEXT: return %[[RESULT]] : i1
+func.func @ceildivui_i1(%arg0: i1, %arg1: i1) -> i1 {
+ %res = arith.ceildivui %arg0, %arg1 : i1
+ return %res : i1
+}
+
+// -----
+
+// CHECK-LABEL: func.func @ceildivi(
+// CHECK-SAME: %[[LHS:.*]]: i32, %[[RHS:.*]]: i32) -> i32 {
+// CHECK-NOT: arith.ceildivsi
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divsi %[[LHS]], %[[RHS]] : i32
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : i32
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : i32
+// CHECK-NEXT: %[[SIGNED_LT:.*]] = arith.cmpi slt, %[[LHS]], %[[RHS]] : i32
+// CHECK-NEXT: %[[UNSIGNED_LT:.*]] = arith.cmpi ult, %[[LHS]], %[[RHS]] : i32
+// CHECK-NEXT: %[[SAME_SIGN:.*]] = arith.cmpi eq, %[[SIGNED_LT]], %[[UNSIGNED_LT]] : i1
+// CHECK-NEXT: %[[ROUND:.*]] = arith.andi %[[INEXACT]], %[[SAME_SIGN]] : i1
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.extui %[[ROUND]] : i1 to i32
+// CHECK-NEXT: %[[RESULT:.*]] = arith.addi %[[Q]], %[[ADJUSTMENT]] : i32
+// CHECK-NEXT: return %[[RESULT]] : i32
+func.func @ceildivi(%arg0: i32, %arg1: i32) -> i32 {
+ %res = arith.ceildivsi %arg0, %arg1 : i32
return %res : i32
-// CHECK: [[ZERO:%.+]] = arith.constant 0 : i32
-// CHECK: [[ISZERO:%.+]] = arith.cmpi eq, %arg0, [[ZERO]] : i32
-// CHECK: [[ONE:%.+]] = arith.constant 1 : i32
-// CHECK: [[SUB:%.+]] = arith.subi %arg0, [[ONE]] : i32
-// CHECK: [[DIV:%.+]] = arith.divui [[SUB]], %arg1 : i32
-// CHECK: [[REM:%.+]] = arith.addi [[DIV]], [[ONE]] : i32
-// CHECK: [[RES:%.+]] = arith.select [[ISZERO]], [[ZERO]], [[REM]] : i32
}
// -----
-// Test unsigned ceil divide with index
-// CHECK-LABEL: func @ceildivui_index
-// CHECK-SAME: ([[ARG0:%.+]]: index, [[ARG1:%.+]]: index) -> index {
-func.func @ceildivui_index(%arg0: index, %arg1: index) -> (index) {
- %res = arith.ceildivui %arg0, %arg1 : index
+// CHECK-LABEL: func.func @ceildivi_index(
+// CHECK-SAME: %[[LHS:.*]]: index, %[[RHS:.*]]: index) -> index {
+// CHECK-NOT: arith.ceildivsi
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.extui
+// CHECK-NEXT: %[[Q:.*]] = arith.divsi %[[LHS]], %[[RHS]] : index
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : index
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : index
+// CHECK-NEXT: %[[SIGNED_LT:.*]] = arith.cmpi slt, %[[LHS]], %[[RHS]] : index
+// CHECK-NEXT: %[[UNSIGNED_LT:.*]] = arith.cmpi ult, %[[LHS]], %[[RHS]] : index
+// CHECK-NEXT: %[[SAME_SIGN:.*]] = arith.cmpi eq, %[[SIGNED_LT]], %[[UNSIGNED_LT]] : i1
+// CHECK-NEXT: %[[ROUND:.*]] = arith.andi %[[INEXACT]], %[[SAME_SIGN]] : i1
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.index_castui %[[ROUND]] : i1 to index
+// CHECK-NEXT: %[[RESULT:.*]] = arith.addi %[[Q]], %[[ADJUSTMENT]] : index
+// CHECK-NEXT: return %[[RESULT]] : index
+func.func @ceildivi_index(%arg0: index, %arg1: index) -> index {
+ %res = arith.ceildivsi %arg0, %arg1 : index
return %res : index
-// CHECK: [[ZERO:%.+]] = arith.constant 0 : index
-// CHECK: [[ISZERO:%.+]] = arith.cmpi eq, %arg0, [[ZERO]] : index
-// CHECK: [[ONE:%.+]] = arith.constant 1 : index
-// CHECK: [[SUB:%.+]] = arith.subi %arg0, [[ONE]] : index
-// CHECK: [[DIV:%.+]] = arith.divui [[SUB]], %arg1 : index
-// CHECK: [[REM:%.+]] = arith.addi [[DIV]], [[ONE]] : index
-// CHECK: [[RES:%.+]] = arith.select [[ISZERO]], [[ZERO]], [[REM]] : 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> {
+// CHECK-LABEL: func.func @ceildivsi_vec(
+// CHECK-SAME: %[[LHS:.*]]: vector<4xi32>, %[[RHS:.*]]: vector<4xi32>) -> vector<4xi32> {
+// CHECK-NOT: arith.ceildivsi
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divsi %[[LHS]], %[[RHS]] : vector<4xi32>
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : vector<4xi32>
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : vector<4xi32>
+// CHECK-NEXT: %[[SIGNED_LT:.*]] = arith.cmpi slt, %[[LHS]], %[[RHS]] : vector<4xi32>
+// CHECK-NEXT: %[[UNSIGNED_LT:.*]] = arith.cmpi ult, %[[LHS]], %[[RHS]] : vector<4xi32>
+// CHECK-NEXT: %[[SAME_SIGN:.*]] = arith.cmpi eq, %[[SIGNED_LT]], %[[UNSIGNED_LT]] : vector<4xi1>
+// CHECK-NEXT: %[[ROUND:.*]] = arith.andi %[[INEXACT]], %[[SAME_SIGN]] : vector<4xi1>
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.extui %[[ROUND]] : vector<4xi1> to vector<4xi32>
+// CHECK-NEXT: %[[RESULT:.*]] = arith.addi %[[Q]], %[[ADJUSTMENT]] : vector<4xi32>
+// CHECK-NEXT: return %[[RESULT]] : vector<4xi32>
+func.func @ceildivsi_vec(%arg0: vector<4xi32>, %arg1: vector<4xi32>) -> vector<4xi32> {
+ %res = arith.ceildivsi %arg0, %arg1 : vector<4xi32>
+ return %res : vector<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @ceildivsi_static_tensor(
+// CHECK-SAME: %[[LHS:.*]]: tensor<2x3xi32>, %[[RHS:.*]]: tensor<2x3xi32>) -> tensor<2x3xi32> {
+// CHECK-NOT: arith.ceildivsi
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divsi %[[LHS]], %[[RHS]] : tensor<2x3xi32>
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : tensor<2x3xi32>
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : tensor<2x3xi32>
+// CHECK-NEXT: %[[SIGNED_LT:.*]] = arith.cmpi slt, %[[LHS]], %[[RHS]] : tensor<2x3xi32>
+// CHECK-NEXT: %[[UNSIGNED_LT:.*]] = arith.cmpi ult, %[[LHS]], %[[RHS]] : tensor<2x3xi32>
+// CHECK-NEXT: %[[SAME_SIGN:.*]] = arith.cmpi eq, %[[SIGNED_LT]], %[[UNSIGNED_LT]] : tensor<2x3xi1>
+// CHECK-NEXT: %[[ROUND:.*]] = arith.andi %[[INEXACT]], %[[SAME_SIGN]] : tensor<2x3xi1>
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.extui %[[ROUND]] : tensor<2x3xi1> to tensor<2x3xi32>
+// CHECK-NEXT: %[[RESULT:.*]] = arith.addi %[[Q]], %[[ADJUSTMENT]] : tensor<2x3xi32>
+// CHECK-NEXT: return %[[RESULT]] : tensor<2x3xi32>
+func.func @ceildivsi_static_tensor(%arg0: tensor<2x3xi32>, %arg1: tensor<2x3xi32>) -> tensor<2x3xi32> {
+ %res = arith.ceildivsi %arg0, %arg1 : tensor<2x3xi32>
+ return %res : tensor<2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @ceildivsi_dynamic_tensor(
+// CHECK-SAME: %[[LHS:.*]]: tensor<8x?xi64>, %[[RHS:.*]]: tensor<8x?xi64>) -> tensor<8x?xi64> {
+// CHECK-NOT: arith.ceildivsi
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divsi %[[LHS]], %[[RHS]] : tensor<8x?xi64>
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : tensor<8x?xi64>
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : tensor<8x?xi64>
+// CHECK-NEXT: %[[SIGNED_LT:.*]] = arith.cmpi slt, %[[LHS]], %[[RHS]] : tensor<8x?xi64>
+// CHECK-NEXT: %[[UNSIGNED_LT:.*]] = arith.cmpi ult, %[[LHS]], %[[RHS]] : tensor<8x?xi64>
+// CHECK-NEXT: %[[SAME_SIGN:.*]] = arith.cmpi eq, %[[SIGNED_LT]], %[[UNSIGNED_LT]] : tensor<8x?xi1>
+// CHECK-NEXT: %[[ROUND:.*]] = arith.andi %[[INEXACT]], %[[SAME_SIGN]] : tensor<8x?xi1>
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.extui %[[ROUND]] : tensor<8x?xi1> to tensor<8x?xi64>
+// CHECK-NEXT: %[[RESULT:.*]] = arith.addi %[[Q]], %[[ADJUSTMENT]] : tensor<8x?xi64>
+// CHECK-NEXT: return %[[RESULT]] : 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> {
+// CHECK-LABEL: func.func @ceildivsi_i1(
+// CHECK-SAME: %[[LHS:.*]]: i1, %[[RHS:.*]]: i1) -> i1 {
+// CHECK-NOT: arith.ceildivsi
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.extui
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divsi %[[LHS]], %[[RHS]] : i1
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : i1
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : i1
+// CHECK-NEXT: %[[SIGNED_LT:.*]] = arith.cmpi slt, %[[LHS]], %[[RHS]] : i1
+// CHECK-NEXT: %[[UNSIGNED_LT:.*]] = arith.cmpi ult, %[[LHS]], %[[RHS]] : i1
+// CHECK-NEXT: %[[SAME_SIGN:.*]] = arith.cmpi eq, %[[SIGNED_LT]], %[[UNSIGNED_LT]] : i1
+// CHECK-NEXT: %[[ROUND:.*]] = arith.andi %[[INEXACT]], %[[SAME_SIGN]] : i1
+// CHECK-NEXT: %[[RESULT:.*]] = arith.addi %[[Q]], %[[ROUND]] : i1
+// CHECK-NEXT: return %[[RESULT]] : i1
+func.func @ceildivsi_i1(%arg0: i1, %arg1: i1) -> i1 {
+ %res = arith.ceildivsi %arg0, %arg1 : i1
+ return %res : i1
+}
+
+// -----
+
+// CHECK-LABEL: func.func @floordivi(
+// CHECK-SAME: %[[LHS:.*]]: i32, %[[RHS:.*]]: i32) -> i32 {
+// CHECK-NOT: arith.floordivsi
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divsi %[[LHS]], %[[RHS]] : i32
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : i32
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : i32
+// CHECK-NEXT: %[[SIGNED_LT:.*]] = arith.cmpi slt, %[[LHS]], %[[RHS]] : i32
+// CHECK-NEXT: %[[UNSIGNED_LT:.*]] = arith.cmpi ult, %[[LHS]], %[[RHS]] : i32
+// CHECK-NEXT: %[[OPPOSITE_SIGN:.*]] = arith.cmpi ne, %[[SIGNED_LT]], %[[UNSIGNED_LT]] : i1
+// CHECK-NEXT: %[[ROUND:.*]] = arith.andi %[[INEXACT]], %[[OPPOSITE_SIGN]] : i1
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.extui %[[ROUND]] : i1 to i32
+// CHECK-NEXT: %[[RESULT:.*]] = arith.subi %[[Q]], %[[ADJUSTMENT]] : i32
+// CHECK-NEXT: return %[[RESULT]] : i32
+func.func @floordivi(%arg0: i32, %arg1: i32) -> i32 {
+ %res = arith.floordivsi %arg0, %arg1 : i32
+ return %res : i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @floordivi_index(
+// CHECK-SAME: %[[LHS:.*]]: index, %[[RHS:.*]]: index) -> index {
+// CHECK-NOT: arith.floordivsi
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.extui
+// CHECK-NEXT: %[[Q:.*]] = arith.divsi %[[LHS]], %[[RHS]] : index
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : index
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : index
+// CHECK-NEXT: %[[SIGNED_LT:.*]] = arith.cmpi slt, %[[LHS]], %[[RHS]] : index
+// CHECK-NEXT: %[[UNSIGNED_LT:.*]] = arith.cmpi ult, %[[LHS]], %[[RHS]] : index
+// CHECK-NEXT: %[[OPPOSITE_SIGN:.*]] = arith.cmpi ne, %[[SIGNED_LT]], %[[UNSIGNED_LT]] : i1
+// CHECK-NEXT: %[[ROUND:.*]] = arith.andi %[[INEXACT]], %[[OPPOSITE_SIGN]] : i1
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.index_castui %[[ROUND]] : i1 to index
+// CHECK-NEXT: %[[RESULT:.*]] = arith.subi %[[Q]], %[[ADJUSTMENT]] : index
+// CHECK-NEXT: return %[[RESULT]] : index
+func.func @floordivi_index(%arg0: index, %arg1: index) -> index {
+ %res = arith.floordivsi %arg0, %arg1 : index
+ return %res : index
+}
+
+// -----
+
+// CHECK-LABEL: func.func @floordivi_vec(
+// CHECK-SAME: %[[LHS:.*]]: vector<4xi32>, %[[RHS:.*]]: vector<4xi32>) -> vector<4xi32> {
+// CHECK-NOT: arith.floordivsi
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divsi %[[LHS]], %[[RHS]] : vector<4xi32>
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : vector<4xi32>
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : vector<4xi32>
+// CHECK-NEXT: %[[SIGNED_LT:.*]] = arith.cmpi slt, %[[LHS]], %[[RHS]] : vector<4xi32>
+// CHECK-NEXT: %[[UNSIGNED_LT:.*]] = arith.cmpi ult, %[[LHS]], %[[RHS]] : vector<4xi32>
+// CHECK-NEXT: %[[OPPOSITE_SIGN:.*]] = arith.cmpi ne, %[[SIGNED_LT]], %[[UNSIGNED_LT]] : vector<4xi1>
+// CHECK-NEXT: %[[ROUND:.*]] = arith.andi %[[INEXACT]], %[[OPPOSITE_SIGN]] : vector<4xi1>
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.extui %[[ROUND]] : vector<4xi1> to vector<4xi32>
+// CHECK-NEXT: %[[RESULT:.*]] = arith.subi %[[Q]], %[[ADJUSTMENT]] : vector<4xi32>
+// CHECK-NEXT: return %[[RESULT]] : vector<4xi32>
+func.func @floordivi_vec(%arg0: vector<4xi32>, %arg1: vector<4xi32>) -> vector<4xi32> {
+ %res = arith.floordivsi %arg0, %arg1 : vector<4xi32>
+ return %res : vector<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @floordivsi_static_tensor(
+// CHECK-SAME: %[[LHS:.*]]: tensor<2x3xi32>, %[[RHS:.*]]: tensor<2x3xi32>) -> tensor<2x3xi32> {
+// CHECK-NOT: arith.floordivsi
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divsi %[[LHS]], %[[RHS]] : tensor<2x3xi32>
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : tensor<2x3xi32>
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : tensor<2x3xi32>
+// CHECK-NEXT: %[[SIGNED_LT:.*]] = arith.cmpi slt, %[[LHS]], %[[RHS]] : tensor<2x3xi32>
+// CHECK-NEXT: %[[UNSIGNED_LT:.*]] = arith.cmpi ult, %[[LHS]], %[[RHS]] : tensor<2x3xi32>
+// CHECK-NEXT: %[[OPPOSITE_SIGN:.*]] = arith.cmpi ne, %[[SIGNED_LT]], %[[UNSIGNED_LT]] : tensor<2x3xi1>
+// CHECK-NEXT: %[[ROUND:.*]] = arith.andi %[[INEXACT]], %[[OPPOSITE_SIGN]] : tensor<2x3xi1>
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.extui %[[ROUND]] : tensor<2x3xi1> to tensor<2x3xi32>
+// CHECK-NEXT: %[[RESULT:.*]] = arith.subi %[[Q]], %[[ADJUSTMENT]] : tensor<2x3xi32>
+// CHECK-NEXT: return %[[RESULT]] : tensor<2x3xi32>
+func.func @floordivsi_static_tensor(%arg0: tensor<2x3xi32>, %arg1: tensor<2x3xi32>) -> tensor<2x3xi32> {
+ %res = arith.floordivsi %arg0, %arg1 : tensor<2x3xi32>
+ return %res : tensor<2x3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @floordivsi_dynamic_tensor(
+// CHECK-SAME: %[[LHS:.*]]: tensor<?x4xi64>, %[[RHS:.*]]: tensor<?x4xi64>) -> tensor<?x4xi64> {
+// CHECK-NOT: arith.floordivsi
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divsi %[[LHS]], %[[RHS]] : tensor<?x4xi64>
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : tensor<?x4xi64>
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : tensor<?x4xi64>
+// CHECK-NEXT: %[[SIGNED_LT:.*]] = arith.cmpi slt, %[[LHS]], %[[RHS]] : tensor<?x4xi64>
+// CHECK-NEXT: %[[UNSIGNED_LT:.*]] = arith.cmpi ult, %[[LHS]], %[[RHS]] : tensor<?x4xi64>
+// CHECK-NEXT: %[[OPPOSITE_SIGN:.*]] = arith.cmpi ne, %[[SIGNED_LT]], %[[UNSIGNED_LT]] : tensor<?x4xi1>
+// CHECK-NEXT: %[[ROUND:.*]] = arith.andi %[[INEXACT]], %[[OPPOSITE_SIGN]] : tensor<?x4xi1>
+// CHECK-NEXT: %[[ADJUSTMENT:.*]] = arith.extui %[[ROUND]] : tensor<?x4xi1> to tensor<?x4xi64>
+// CHECK-NEXT: %[[RESULT:.*]] = arith.subi %[[Q]], %[[ADJUSTMENT]] : tensor<?x4xi64>
+// CHECK-NEXT: return %[[RESULT]] : 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.func @floordivsi_i1(
+// CHECK-SAME: %[[LHS:.*]]: i1, %[[RHS:.*]]: i1) -> i1 {
+// CHECK-NOT: arith.floordivsi
+// CHECK-NOT: arith.constant
+// CHECK-NOT: tensor.
+// CHECK-NOT: arith.extui
+// CHECK-NOT: arith.index_castui
+// CHECK-NEXT: %[[Q:.*]] = arith.divsi %[[LHS]], %[[RHS]] : i1
+// CHECK-NEXT: %[[PRODUCT:.*]] = arith.muli %[[Q]], %[[RHS]] : i1
+// CHECK-NEXT: %[[INEXACT:.*]] = arith.cmpi ne, %[[PRODUCT]], %[[LHS]] : i1
+// CHECK-NEXT: %[[SIGNED_LT:.*]] = arith.cmpi slt, %[[LHS]], %[[RHS]] : i1
+// CHECK-NEXT: %[[UNSIGNED_LT:.*]] = arith.cmpi ult, %[[LHS]], %[[RHS]] : i1
+// CHECK-NEXT: %[[OPPOSITE_SIGN:.*]] = arith.cmpi ne, %[[SIGNED_LT]], %[[UNSIGNED_LT]] : i1
+// CHECK-NEXT: %[[ROUND:.*]] = arith.andi %[[INEXACT]], %[[OPPOSITE_SIGN]] : i1
+// CHECK-NEXT: %[[RESULT:.*]] = arith.subi %[[Q]], %[[ROUND]] : i1
+// CHECK-NEXT: return %[[RESULT]] : i1
+func.func @floordivsi_i1(%arg0: i1, %arg1: i1) -> i1 {
+ %res = arith.floordivsi %arg0, %arg1 : i1
+ return %res : i1
+}
+
+// -----
+
+// VALUES-LABEL: func.func @ceildivui_values()
+// VALUES-DAG: %[[THREE:.*]] = arith.constant 3 : i8
+// VALUES-DAG: %[[TWO:.*]] = arith.constant 2 : i8
+// VALUES-DAG: %[[ZERO:.*]] = arith.constant 0 : i8
+// VALUES: return %[[THREE]], %[[TWO]], %[[ZERO]] : i8, i8, i8
+func.func @ceildivui_values() -> (i8, i8, i8) {
+ %zero = arith.constant 0 : i8
+ %two = arith.constant 2 : i8
+ %four = arith.constant 4 : i8
+ %five = arith.constant 5 : i8
+ %inexact = arith.ceildivui %five, %two : i8
+ %exact = arith.ceildivui %four, %two : i8
+ %zeroDividend = arith.ceildivui %zero, %two : i8
+ return %inexact, %exact, %zeroDividend : i8, i8, i8
+}
+
+// VALUES-LABEL: func.func @ceildivsi_values()
+// VALUES-DAG: %[[THREE:.*]] = arith.constant 3 : i8
+// VALUES-DAG: %[[NEG_TWO:.*]] = arith.constant -2 : i8
+// VALUES: return %[[THREE]], %[[NEG_TWO]], %[[NEG_TWO]], %[[THREE]] : i8, i8, i8, i8
+func.func @ceildivsi_values() -> (i8, i8, i8, i8) {
+ %negFive = arith.constant -5 : i8
+ %negTwo = arith.constant -2 : i8
+ %two = arith.constant 2 : i8
+ %five = arith.constant 5 : i8
+ %posPos = arith.ceildivsi %five, %two : i8
+ %negPos = arith.ceildivsi %negFive, %two : i8
+ %posNeg = arith.ceildivsi %five, %negTwo : i8
+ %negNeg = arith.ceildivsi %negFive, %negTwo : i8
+ return %posPos, %negPos, %posNeg, %negNeg : i8, i8, i8, i8
+}
+
+// VALUES-LABEL: func.func @floordivsi_values()
+// VALUES-DAG: %[[TWO:.*]] = arith.constant 2 : i8
+// VALUES-DAG: %[[NEG_THREE:.*]] = arith.constant -3 : i8
+// VALUES: return %[[TWO]], %[[NEG_THREE]], %[[NEG_THREE]], %[[TWO]] : i8, i8, i8, i8
+func.func @floordivsi_values() -> (i8, i8, i8, i8) {
+ %negFive = arith.constant -5 : i8
+ %negTwo = arith.constant -2 : i8
+ %two = arith.constant 2 : i8
+ %five = arith.constant 5 : i8
+ %posPos = arith.floordivsi %five, %two : i8
+ %negPos = arith.floordivsi %negFive, %two : i8
+ %posNeg = arith.floordivsi %five, %negTwo : i8
+ %negNeg = arith.floordivsi %negFive, %negTwo : i8
+ return %posPos, %negPos, %posNeg, %negNeg : i8, i8, i8, i8
}
// -----
>From e05cd53e02b5ab91d98317901d835c37c1b2070b Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Tue, 14 Jul 2026 16:50:11 +0800
Subject: [PATCH 6/8] [MLIR] Update division consumer checks
---
.../Conversion/ArithToLLVM/arith-to-llvm.mlir | 52 ++++++++-----------
.../test/Conversion/ConvertToSPIRV/arith.mlir | 51 ++++++++++--------
2 files changed, 52 insertions(+), 51 deletions(-)
diff --git a/mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir b/mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir
index 43f21561e6544..133c10eeb6ed8 100644
--- a/mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir
+++ b/mlir/test/Conversion/ArithToLLVM/arith-to-llvm.mlir
@@ -764,17 +764,15 @@ func.func @select_complex(%arg0 : i1, %arg1 : complex<f32>, %arg2 : complex<f32>
// CHECK-LABEL: @ceildivsi
// CHECK-SAME: %[[ARG0:.*]]: i64, %[[ARG1:.*]]: i64) -> i64
func.func @ceildivsi(%arg0 : i64, %arg1 : i64) -> i64 {
- // CHECK: %[[ZERO:.+]] = llvm.mlir.constant(0 : i64) : i64
- // CHECK: %[[ONE:.+]] = llvm.mlir.constant(1 : i64) : i64
- // CHECK: %[[DIV:.+]] = llvm.sdiv %[[ARG0]], %[[ARG1]] : i64
- // CHECK: %[[MUL:.+]] = llvm.mul %[[DIV]], %[[ARG1]] : i64
- // CHECK: %[[NEXACT:.+]] = llvm.icmp "ne" %[[ARG0]], %[[MUL]] : i64
- // CHECK: %[[NNEG:.+]] = llvm.icmp "slt" %[[ARG0]], %[[ZERO]] : i64
- // CHECK: %[[MNEG:.+]] = llvm.icmp "slt" %[[ARG1]], %[[ZERO]] : i64
- // CHECK: %[[SAMESIGN:.+]] = llvm.icmp "eq" %[[NNEG]], %[[MNEG]] : i1
- // CHECK: %[[SHOULDROUND:.+]] = llvm.and %[[NEXACT]], %[[SAMESIGN]] : i1
- // CHECK: %[[CEIL:.+]] = llvm.add %[[DIV]], %[[ONE]] : i64
- // CHECK: %[[RES:.+]] = llvm.select %[[SHOULDROUND]], %[[CEIL]], %[[DIV]] : i1, i64
+ // CHECK: %[[Q:.+]] = llvm.sdiv %[[ARG0]], %[[ARG1]] : i64
+ // CHECK: %[[PRODUCT:.+]] = llvm.mul %[[Q]], %[[ARG1]] : i64
+ // CHECK: %[[INEXACT:.+]] = llvm.icmp "ne" %[[PRODUCT]], %[[ARG0]] : i64
+ // CHECK: %[[SIGNED_LT:.+]] = llvm.icmp "slt" %[[ARG0]], %[[ARG1]] : i64
+ // CHECK: %[[UNSIGNED_LT:.+]] = llvm.icmp "ult" %[[ARG0]], %[[ARG1]] : i64
+ // CHECK: %[[SAME_SIGN:.+]] = llvm.icmp "eq" %[[SIGNED_LT]], %[[UNSIGNED_LT]] : i1
+ // CHECK: %[[ROUND:.+]] = llvm.and %[[INEXACT]], %[[SAME_SIGN]] : i1
+ // CHECK: %[[ADJUSTMENT:.+]] = llvm.zext %[[ROUND]] : i1 to i64
+ // CHECK: %[[RES:.+]] = llvm.add %[[Q]], %[[ADJUSTMENT]] : i64
%0 = arith.ceildivsi %arg0, %arg1 : i64
return %0: i64
}
@@ -782,13 +780,11 @@ func.func @ceildivsi(%arg0 : i64, %arg1 : i64) -> i64 {
// CHECK-LABEL: @ceildivui
// CHECK-SAME: %[[ARG0:.*]]: i32) -> i32
func.func @ceildivui(%arg0 : i32) -> i32 {
-// CHECK: %[[CST0:.*]] = llvm.mlir.constant(0 : i32) : i32
-// CHECK: %[[CMP0:.*]] = llvm.icmp "eq" %[[ARG0]], %[[CST0]] : i32
-// CHECK: %[[CST1:.*]] = llvm.mlir.constant(1 : i32) : i32
-// CHECK: %[[SUB0:.*]] = llvm.sub %[[ARG0]], %[[CST1]] : i32
-// CHECK: %[[DIV0:.*]] = llvm.udiv %[[SUB0]], %[[ARG0]] : i32
-// CHECK: %[[ADD0:.*]] = llvm.add %[[DIV0]], %[[CST1]] : i32
-// CHECK: %[[SEL0:.*]] = llvm.select %[[CMP0]], %[[CST0]], %[[ADD0]] : i1, i32
+// CHECK: %[[Q:.*]] = llvm.udiv %[[ARG0]], %[[ARG0]] : i32
+// CHECK: %[[PRODUCT:.*]] = llvm.mul %[[Q]], %[[ARG0]] : i32
+// CHECK: %[[INEXACT:.*]] = llvm.icmp "ne" %[[PRODUCT]], %[[ARG0]] : i32
+// CHECK: %[[ADJUSTMENT:.*]] = llvm.zext %[[INEXACT]] : i1 to i32
+// CHECK: %[[RES:.*]] = llvm.add %[[Q]], %[[ADJUSTMENT]] : i32
%0 = arith.ceildivui %arg0, %arg0 : i32
return %0: i32
}
@@ -798,17 +794,15 @@ func.func @ceildivui(%arg0 : i32) -> i32 {
// CHECK-LABEL: @floordivsi
// CHECK-SAME: %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32) -> i32
func.func @floordivsi(%arg0 : i32, %arg1 : i32) -> i32 {
- // CHECK: %[[SDIV:.*]] = llvm.sdiv %[[ARG0]], %[[ARG1]] : i32
- // CHECK: %[[MUL0:.*]] = llvm.mul %[[SDIV]], %[[ARG1]] : i32
- // CHECK: %[[CMP0:.*]] = llvm.icmp "ne" %[[ARG0]], %[[MUL0]] : i32
- // CHECK: %[[CST0:.*]] = llvm.mlir.constant(0 : i32) : i32
- // CHECK: %[[CMP1:.*]] = llvm.icmp "slt" %[[ARG0]], %[[CST0]] : i32
- // CHECK: %[[CMP2:.*]] = llvm.icmp "slt" %[[ARG1]], %[[CST0]] : i32
- // CHECK: %[[CMP3:.*]] = llvm.icmp "ne" %[[CMP1]], %[[CMP2]] : i1
- // CHECK: %[[AND:.*]] = llvm.and %[[CMP0]], %[[CMP3]] : i1
- // CHECK: %[[CST1:.*]] = llvm.mlir.constant(-1 : i32) : i32
- // CHECK: %[[ADD:.*]] = llvm.add %[[SDIV]], %[[CST1]] : i32
- // CHECK: %[[SEL:.*]] = llvm.select %[[AND]], %[[ADD]], %[[SDIV]] : i1, i32
+ // CHECK: %[[Q:.*]] = llvm.sdiv %[[ARG0]], %[[ARG1]] : i32
+ // CHECK: %[[PRODUCT:.*]] = llvm.mul %[[Q]], %[[ARG1]] : i32
+ // CHECK: %[[INEXACT:.*]] = llvm.icmp "ne" %[[PRODUCT]], %[[ARG0]] : i32
+ // CHECK: %[[SIGNED_LT:.*]] = llvm.icmp "slt" %[[ARG0]], %[[ARG1]] : i32
+ // CHECK: %[[UNSIGNED_LT:.*]] = llvm.icmp "ult" %[[ARG0]], %[[ARG1]] : i32
+ // CHECK: %[[OPPOSITE_SIGN:.*]] = llvm.icmp "ne" %[[SIGNED_LT]], %[[UNSIGNED_LT]] : i1
+ // CHECK: %[[ROUND:.*]] = llvm.and %[[INEXACT]], %[[OPPOSITE_SIGN]] : i1
+ // CHECK: %[[ADJUSTMENT:.*]] = llvm.zext %[[ROUND]] : i1 to i32
+ // CHECK: %[[RES:.*]] = llvm.sub %[[Q]], %[[ADJUSTMENT]] : i32
%0 = arith.floordivsi %arg0, %arg1 : i32
return %0 : i32
}
diff --git a/mlir/test/Conversion/ConvertToSPIRV/arith.mlir b/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
index fa204c4df8ace..58e5ac10d3a1b 100644
--- a/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
+++ b/mlir/test/Conversion/ConvertToSPIRV/arith.mlir
@@ -37,11 +37,13 @@ func.func @int32_scalar_srem(%lhs: i32, %rhs: i32) {
// CHECK-LABEL: @scalar_ceildivui
// CHECK-SAME: (%[[LHS:.+]]: i32, %[[RHS:.+]]: i32)
func.func @scalar_ceildivui(%lhs: i32, %rhs: i32) -> i32 {
- // CHECK: %[[ISZERO:.+]] = spirv.IEqual %[[LHS]], %{{.+}} : i32
- // CHECK: %[[MINUSONE:.+]] = spirv.ISub %[[LHS]], %{{.+}} : i32
- // CHECK: %[[Q:.+]] = spirv.UDiv %[[MINUSONE]], %[[RHS]] : i32
- // CHECK: %[[PLUSONE:.+]] = spirv.IAdd %[[Q]], %{{.+}} : i32
- // CHECK: %[[R:.+]] = spirv.Select %[[ISZERO]], %{{.+}}, %[[PLUSONE]] : i1, i32
+ // CHECK: %[[Q:.+]] = spirv.UDiv %[[LHS]], %[[RHS]] : i32
+ // CHECK: %[[PRODUCT:.+]] = spirv.IMul %[[Q]], %[[RHS]] : i32
+ // CHECK: %[[INEXACT:.+]] = spirv.INotEqual %[[PRODUCT]], %[[LHS]] : i32
+ // CHECK: %[[ZERO:.+]] = spirv.Constant 0 : i32
+ // CHECK: %[[ONE:.+]] = spirv.Constant 1 : i32
+ // CHECK: %[[ADJUSTMENT:.+]] = spirv.Select %[[INEXACT]], %[[ONE]], %[[ZERO]] : i1, i32
+ // CHECK: %[[R:.+]] = spirv.IAdd %[[Q]], %[[ADJUSTMENT]] : i32
// CHECK: spirv.ReturnValue %[[R]]
%0 = arith.ceildivui %lhs, %rhs : i32
return %0 : i32
@@ -51,14 +53,16 @@ func.func @scalar_ceildivui(%lhs: i32, %rhs: i32) -> i32 {
// CHECK-SAME: (%[[LHS:.+]]: i32, %[[RHS:.+]]: i32)
func.func @scalar_ceildivsi(%lhs: i32, %rhs: i32) -> i32 {
// CHECK: %[[Q:.+]] = spirv.SDiv %[[LHS]], %[[RHS]] : i32
- // CHECK: %[[PROD:.+]] = spirv.IMul %[[Q]], %[[RHS]] : i32
- // CHECK: %[[NE:.+]] = spirv.INotEqual %[[LHS]], %[[PROD]] : i32
- // CHECK: %[[LNEG:.+]] = spirv.SLessThan %[[LHS]], %{{.+}} : i32
- // CHECK: %[[RNEG:.+]] = spirv.SLessThan %[[RHS]], %{{.+}} : i32
- // CHECK: %[[SAMESIGN:.+]] = spirv.LogicalEqual %[[LNEG]], %[[RNEG]] : i1
- // CHECK: %[[COND:.+]] = spirv.LogicalAnd %[[NE]], %[[SAMESIGN]] : i1
- // CHECK: %[[QP1:.+]] = spirv.IAdd %[[Q]], %{{.+}} : i32
- // CHECK: %[[R:.+]] = spirv.Select %[[COND]], %[[QP1]], %[[Q]] : i1, i32
+ // CHECK: %[[PRODUCT:.+]] = spirv.IMul %[[Q]], %[[RHS]] : i32
+ // CHECK: %[[INEXACT:.+]] = spirv.INotEqual %[[PRODUCT]], %[[LHS]] : i32
+ // CHECK: %[[SIGNED_LT:.+]] = spirv.SLessThan %[[LHS]], %[[RHS]] : i32
+ // CHECK: %[[UNSIGNED_LT:.+]] = spirv.ULessThan %[[LHS]], %[[RHS]] : i32
+ // CHECK: %[[SAME_SIGN:.+]] = spirv.LogicalEqual %[[SIGNED_LT]], %[[UNSIGNED_LT]] : i1
+ // CHECK: %[[ROUND:.+]] = spirv.LogicalAnd %[[INEXACT]], %[[SAME_SIGN]] : i1
+ // CHECK: %[[ZERO:.+]] = spirv.Constant 0 : i32
+ // CHECK: %[[ONE:.+]] = spirv.Constant 1 : i32
+ // CHECK: %[[ADJUSTMENT:.+]] = spirv.Select %[[ROUND]], %[[ONE]], %[[ZERO]] : i1, i32
+ // CHECK: %[[R:.+]] = spirv.IAdd %[[Q]], %[[ADJUSTMENT]] : i32
// CHECK: spirv.ReturnValue %[[R]]
%0 = arith.ceildivsi %lhs, %rhs : i32
return %0 : i32
@@ -68,14 +72,16 @@ func.func @scalar_ceildivsi(%lhs: i32, %rhs: i32) -> i32 {
// CHECK-SAME: (%[[LHS:.+]]: i32, %[[RHS:.+]]: i32)
func.func @scalar_floordivsi(%lhs: i32, %rhs: i32) -> i32 {
// CHECK: %[[Q:.+]] = spirv.SDiv %[[LHS]], %[[RHS]] : i32
- // CHECK: %[[PROD:.+]] = spirv.IMul %[[Q]], %[[RHS]] : i32
- // CHECK: %[[NE:.+]] = spirv.INotEqual %[[LHS]], %[[PROD]] : i32
- // CHECK: %[[LNEG:.+]] = spirv.SLessThan %[[LHS]], %{{.+}} : i32
- // CHECK: %[[RNEG:.+]] = spirv.SLessThan %[[RHS]], %{{.+}} : i32
- // CHECK: %[[DIFFSIGN:.+]] = spirv.LogicalNotEqual %[[LNEG]], %[[RNEG]] : i1
- // CHECK: %[[COND:.+]] = spirv.LogicalAnd %[[NE]], %[[DIFFSIGN]] : i1
- // CHECK: %[[QM1:.+]] = spirv.IAdd %[[Q]], %{{.+}} : i32
- // CHECK: %[[R:.+]] = spirv.Select %[[COND]], %[[QM1]], %[[Q]] : i1, i32
+ // CHECK: %[[PRODUCT:.+]] = spirv.IMul %[[Q]], %[[RHS]] : i32
+ // CHECK: %[[INEXACT:.+]] = spirv.INotEqual %[[PRODUCT]], %[[LHS]] : i32
+ // CHECK: %[[SIGNED_LT:.+]] = spirv.SLessThan %[[LHS]], %[[RHS]] : i32
+ // CHECK: %[[UNSIGNED_LT:.+]] = spirv.ULessThan %[[LHS]], %[[RHS]] : i32
+ // CHECK: %[[OPPOSITE_SIGN:.+]] = spirv.LogicalNotEqual %[[SIGNED_LT]], %[[UNSIGNED_LT]] : i1
+ // CHECK: %[[ROUND:.+]] = spirv.LogicalAnd %[[INEXACT]], %[[OPPOSITE_SIGN]] : i1
+ // CHECK: %[[ZERO:.+]] = spirv.Constant 0 : i32
+ // CHECK: %[[ONE:.+]] = spirv.Constant 1 : i32
+ // CHECK: %[[ADJUSTMENT:.+]] = spirv.Select %[[ROUND]], %[[ONE]], %[[ZERO]] : i1, i32
+ // CHECK: %[[R:.+]] = spirv.ISub %[[Q]], %[[ADJUSTMENT]] : i32
// CHECK: spirv.ReturnValue %[[R]]
%0 = arith.floordivsi %lhs, %rhs : i32
return %0 : i32
@@ -84,7 +90,8 @@ func.func @scalar_floordivsi(%lhs: i32, %rhs: i32) -> i32 {
// CHECK-LABEL: @vector_ceildivsi
func.func @vector_ceildivsi(%lhs: vector<4xi32>, %rhs: vector<4xi32>) -> vector<4xi32> {
// CHECK: spirv.SDiv %{{.*}}, %{{.*}} : vector<4xi32>
- // CHECK: spirv.Select %{{.*}}, %{{.*}}, %{{.*}} : vector<4xi1>, vector<4xi32>
+ // CHECK: %[[ADJUSTMENT:.+]] = spirv.Select %{{.*}}, %{{.*}}, %{{.*}} : vector<4xi1>, vector<4xi32>
+ // CHECK: spirv.IAdd %{{.*}}, %[[ADJUSTMENT]] : vector<4xi32>
%0 = arith.ceildivsi %lhs, %rhs : vector<4xi32>
return %0 : vector<4xi32>
}
>From 9595a0c00124be170b8512a650dbb0db871d4256 Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Tue, 14 Jul 2026 16:50:27 +0800
Subject: [PATCH 7/8] [MLIR][SPIR-V] Support boolean index casts
---
.../Conversion/ArithToSPIRV/ArithToSPIRV.cpp | 3 +-
.../ArithToSPIRV/arith-to-spirv.mlir | 28 ++++++++++++
.../ArithToSPIRV/ceil-floor-div-index.mlir | 45 +++++++++++++++++++
3 files changed, 75 insertions(+), 1 deletion(-)
create mode 100644 mlir/test/Conversion/ArithToSPIRV/ceil-floor-div-index.mlir
diff --git a/mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp b/mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp
index 9a6d330db72fe..afe8eaab861b0 100644
--- a/mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp
+++ b/mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp
@@ -675,7 +675,7 @@ struct ShRSIBoolPattern final : public OpConversionPattern<arith::ShRSIOp> {
/// Converts an op whose i1 (or vector of i1) source selects between one and
/// zero of the destination type, i.e. spirv.Select(src, one, zero). Shared by
-/// arith.uitofp, arith.extui, and arith.index_cast on boolean sources.
+/// arith.uitofp, arith.extui, and index casts on boolean sources.
template <typename ArithOp>
struct BoolToValuePattern final : public OpConversionPattern<ArithOp> {
using OpConversionPattern<ArithOp>::OpConversionPattern;
@@ -1508,6 +1508,7 @@ void mlir::arith::populateArithToSPIRVPatterns(
TypeCastingOpPattern<arith::IndexCastOp, spirv::SConvertOp>,
IndexCastIndexI1Pattern, BoolToValuePattern<arith::IndexCastOp>,
TypeCastingOpPattern<arith::IndexCastUIOp, spirv::UConvertOp>,
+ BoolToValuePattern<arith::IndexCastUIOp>,
TypeCastingOpPattern<arith::BitcastOp, spirv::BitcastOp>,
CmpIOpBooleanPattern, CmpIOpPattern,
CmpFOpNanNonePattern, CmpFOpPattern,
diff --git a/mlir/test/Conversion/ArithToSPIRV/arith-to-spirv.mlir b/mlir/test/Conversion/ArithToSPIRV/arith-to-spirv.mlir
index b6a488f0dad73..b8dbce76331d8 100644
--- a/mlir/test/Conversion/ArithToSPIRV/arith-to-spirv.mlir
+++ b/mlir/test/Conversion/ArithToSPIRV/arith-to-spirv.mlir
@@ -893,6 +893,34 @@ func.func @index_casti1index_3(%arg0 : vector<3xi1>) {
return
}
+// CHECK-LABEL: index_castuii1index_1
+func.func @index_castuii1index_1(%arg0 : i1) {
+ // CHECK: %[[ZERO:.+]] = spirv.Constant 0 : i32
+ // CHECK: %[[ONE:.+]] = spirv.Constant 1 : i32
+ // CHECK: spirv.Select %{{.+}}, %[[ONE]], %[[ZERO]] : i1, i32
+ %0 = arith.index_castui %arg0 : i1 to index
+ return
+}
+
+// CHECK-LABEL: index_castuii1index_2
+func.func @index_castuii1index_2(%arg0 : vector<1xi1>) -> vector<1xindex> {
+ // Single-element vectors do not exist in SPIRV.
+ // CHECK: %[[ZERO:.+]] = spirv.Constant 0 : i32
+ // CHECK: %[[ONE:.+]] = spirv.Constant 1 : i32
+ // CHECK: spirv.Select %{{.+}}, %[[ONE]], %[[ZERO]] : i1, i32
+ %0 = arith.index_castui %arg0 : vector<1xi1> to vector<1xindex>
+ return %0 : vector<1xindex>
+}
+
+// CHECK-LABEL: index_castuii1index_3
+func.func @index_castuii1index_3(%arg0 : vector<3xi1>) {
+ // CHECK: %[[ZERO:.+]] = spirv.Constant dense<0> : vector<3xi32>
+ // CHECK: %[[ONE:.+]] = spirv.Constant dense<1> : vector<3xi32>
+ // CHECK: spirv.Select %{{.+}}, %[[ONE]], %[[ZERO]] : vector<3xi1>, vector<3xi32>
+ %0 = arith.index_castui %arg0 : vector<3xi1> to vector<3xindex>
+ return
+}
+
// CHECK-LABEL: @bit_cast
func.func @bit_cast(%arg0: vector<2xf32>, %arg1: i64) {
// CHECK: spirv.Bitcast %{{.+}} : vector<2xf32> to vector<2xi32>
diff --git a/mlir/test/Conversion/ArithToSPIRV/ceil-floor-div-index.mlir b/mlir/test/Conversion/ArithToSPIRV/ceil-floor-div-index.mlir
new file mode 100644
index 0000000000000..961127306468a
--- /dev/null
+++ b/mlir/test/Conversion/ArithToSPIRV/ceil-floor-div-index.mlir
@@ -0,0 +1,45 @@
+// RUN: mlir-opt %s -arith-expand -convert-arith-to-spirv | FileCheck %s
+
+// CHECK-LABEL: func.func @ceildivui_index
+// CHECK: %[[Q:.*]] = spirv.UDiv
+// CHECK: %[[PRODUCT:.*]] = spirv.IMul %[[Q]]
+// CHECK: %[[INEXACT:.*]] = spirv.INotEqual %[[PRODUCT]]
+// CHECK: %[[ADJUSTMENT:.*]] = spirv.Select %[[INEXACT]]
+// CHECK: %[[RESULT:.*]] = spirv.IAdd %[[Q]], %[[ADJUSTMENT]]
+// CHECK: return %{{.*}} : index
+func.func @ceildivui_index(%lhs: index, %rhs: index) -> index {
+ %result = arith.ceildivui %lhs, %rhs : index
+ return %result : index
+}
+
+// CHECK-LABEL: func.func @ceildivsi_index
+// CHECK: %[[Q:.*]] = spirv.SDiv
+// CHECK: %[[PRODUCT:.*]] = spirv.IMul %[[Q]]
+// CHECK: %[[INEXACT:.*]] = spirv.INotEqual %[[PRODUCT]]
+// CHECK: %[[SIGNED_LT:.*]] = spirv.SLessThan
+// CHECK: %[[UNSIGNED_LT:.*]] = spirv.ULessThan
+// CHECK: %[[SAME_SIGN:.*]] = spirv.LogicalEqual %[[SIGNED_LT]], %[[UNSIGNED_LT]]
+// CHECK: %[[ROUND:.*]] = spirv.LogicalAnd %[[INEXACT]], %[[SAME_SIGN]]
+// CHECK: %[[ADJUSTMENT:.*]] = spirv.Select %[[ROUND]]
+// CHECK: %[[RESULT:.*]] = spirv.IAdd %[[Q]], %[[ADJUSTMENT]]
+// CHECK: return %{{.*}} : index
+func.func @ceildivsi_index(%lhs: index, %rhs: index) -> index {
+ %result = arith.ceildivsi %lhs, %rhs : index
+ return %result : index
+}
+
+// CHECK-LABEL: func.func @floordivsi_index
+// CHECK: %[[Q:.*]] = spirv.SDiv
+// CHECK: %[[PRODUCT:.*]] = spirv.IMul %[[Q]]
+// CHECK: %[[INEXACT:.*]] = spirv.INotEqual %[[PRODUCT]]
+// CHECK: %[[SIGNED_LT:.*]] = spirv.SLessThan
+// CHECK: %[[UNSIGNED_LT:.*]] = spirv.ULessThan
+// CHECK: %[[OPPOSITE_SIGN:.*]] = spirv.LogicalNotEqual %[[SIGNED_LT]], %[[UNSIGNED_LT]]
+// CHECK: %[[ROUND:.*]] = spirv.LogicalAnd %[[INEXACT]], %[[OPPOSITE_SIGN]]
+// CHECK: %[[ADJUSTMENT:.*]] = spirv.Select %[[ROUND]]
+// CHECK: %[[RESULT:.*]] = spirv.ISub %[[Q]], %[[ADJUSTMENT]]
+// CHECK: return %{{.*}} : index
+func.func @floordivsi_index(%lhs: index, %rhs: index) -> index {
+ %result = arith.floordivsi %lhs, %rhs : index
+ return %result : index
+}
>From 9a358571792d420092cbfb666700a3ff491b86b2 Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Tue, 14 Jul 2026 16:50:48 +0800
Subject: [PATCH 8/8] [MLIR][Affine] Update division rationale
---
mlir/lib/Dialect/Affine/Utils/Utils.cpp | 9 +++------
1 file changed, 3 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 7043083298615..4c4f6981fac74 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -113,12 +113,9 @@ class AffineApplyExpander
/// let quotient = absolute / b in
/// negative ? -quotient - 1 : quotient
///
- /// Note: this lowering does not use arith.floordivsi because the lowering of
- /// that to arith.divsi (see populateCeilFloorDivExpandOpsPatterns) generates
- /// not one but two arith.divsi. That could be changed to one divsi, but one
- /// way or another, going through arith.floordivsi will result in more complex
- /// IR because arith.floordivsi is more general than affine floordiv in that
- /// it supports negative RHS.
+ /// Note: this lowering does not use arith.floordivsi because that operation
+ /// is more general than affine floordiv: it also supports negative RHS. The
+ /// dedicated positive-divisor lowering therefore produces simpler IR.
Value visitFloorDivExpr(AffineBinaryOpExpr expr) {
if (auto rhsConst = dyn_cast<AffineConstantExpr>(expr.getRHS())) {
if (rhsConst.getValue() <= 0) {
More information about the Mlir-commits
mailing list