[Mlir-commits] [mlir] [MLIR][Math] Handle dynamic-shaped tensors without assertion (PR #204510)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jun 26 01:51:53 PDT 2026
https://github.com/Peruere1828 updated https://github.com/llvm/llvm-project/pull/204510
>From fbce9087ed3614643ffc184c7de3f5ef438c5055 Mon Sep 17 00:00:00 2001
From: Peruere1828 <fmyh1828 at gmail.com>
Date: Thu, 18 Jun 2026 05:27:10 +0000
Subject: [PATCH 1/3] Handle dynamic-shaped tensors without assertion
Add `isDynamicShaped` helper to detect shaped types without static
shape (e.g., tensor<?xf32>). Return failure in 10 ops that use
`createFloatConst` (which assumes static shapes for DenseElementsAttr):
sinh, cosh, tanh, asinh, acosh, atanh, exp2, round, roundeven, ctlz
Refactor existing guards in ceil and rsqrt to use the shared helper
for consistency.
Add lit test coverage for each op with both ?-shaped and unranked
tensors, verifying the ops are preserved unchanged.
Fix issue #203753
---
.../lib/Dialect/Math/Transforms/ExpandOps.cpp | 56 +++-
mlir/test/Dialect/Math/expand-math.mlir | 260 ++++++++++++++++++
2 files changed, 307 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
index f76ddfae2a67a..27bcaee5c8303 100644
--- a/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
@@ -25,6 +25,15 @@ namespace mlir::math {
#include "mlir/Dialect/Math/Transforms/Passes.h.inc"
} // namespace mlir::math
+/// Check if the type is a shaped type with dynamic shape
+/// (e.g., tensor<?xf32>). Returns false for scalar types and static-shaped
+/// types, allowing expansion to proceed. Only dynamic-shaped tensors/vectors
+/// are rejected since the generated constant folding assumes static shapes.
+static bool isDynamicShaped(Type type) {
+ auto shapedTy = dyn_cast<ShapedType>(type);
+ return shapedTy && !shapedTy.hasStaticShape();
+}
+
/// Create a float constant.
static Value createFloatConst(Location loc, Type type, APFloat value,
OpBuilder &b) {
@@ -77,6 +86,9 @@ static LogicalResult convertSinhOp(math::SinhOp op, PatternRewriter &rewriter) {
Value operand = op.getOperand();
Type opType = operand.getType();
+ if (isDynamicShaped(opType))
+ return failure();
+
Value exp = math::ExpOp::create(b, operand);
Value neg = arith::NegFOp::create(b, operand);
Value nexp = math::ExpOp::create(b, neg);
@@ -93,6 +105,9 @@ static LogicalResult convertCoshOp(math::CoshOp op, PatternRewriter &rewriter) {
Value operand = op.getOperand();
Type opType = operand.getType();
+ if (isDynamicShaped(opType))
+ return failure();
+
Value exp = math::ExpOp::create(b, operand);
Value neg = arith::NegFOp::create(b, operand);
Value nexp = math::ExpOp::create(b, neg);
@@ -113,6 +128,10 @@ static LogicalResult convertCoshOp(math::CoshOp op, PatternRewriter &rewriter) {
/// result by `sign(x)` to retain sign of the real result.
static LogicalResult convertTanhOp(math::TanhOp op, PatternRewriter &rewriter) {
auto floatType = op.getOperand().getType();
+
+ if (isDynamicShaped(floatType))
+ return failure();
+
Location loc = op.getLoc();
Value zero = createFloatConst(loc, floatType, 0.0, rewriter);
Value one = createFloatConst(loc, floatType, 1.0, rewriter);
@@ -162,6 +181,9 @@ static LogicalResult convertAsinhOp(math::AsinhOp op,
Value operand = op.getOperand();
Type opType = operand.getType();
+ if (isDynamicShaped(opType))
+ return failure();
+
Value one = createFloatConst(op->getLoc(), opType, 1.0, rewriter);
Value fma = math::FmaOp::create(b, operand, operand, one);
Value sqrt = math::SqrtOp::create(b, fma);
@@ -178,6 +200,9 @@ static LogicalResult convertAcoshOp(math::AcoshOp op,
Value operand = op.getOperand();
Type opType = operand.getType();
+ if (isDynamicShaped(opType))
+ return failure();
+
Value negOne = createFloatConst(op->getLoc(), opType, -1.0, rewriter);
Value fma = math::FmaOp::create(b, operand, operand, negOne);
Value sqrt = math::SqrtOp::create(b, fma);
@@ -194,6 +219,9 @@ static LogicalResult convertAtanhOp(math::AtanhOp op,
Value operand = op.getOperand();
Type opType = operand.getType();
+ if (isDynamicShaped(opType))
+ return failure();
+
Value one = createFloatConst(op->getLoc(), opType, 1.0, rewriter);
Value add = arith::AddFOp::create(b, operand, one);
Value neg = arith::NegFOp::create(b, operand);
@@ -224,14 +252,13 @@ static LogicalResult convertFmaFOp(math::FmaOp op, PatternRewriter &rewriter) {
// if (x > y) then incr = 1 else incr = 0
// y = y + incr <= replace this op with the ceilf op.
static LogicalResult convertCeilOp(math::CeilOp op, PatternRewriter &rewriter) {
- // Creating constants assumes the static shaped type.
- auto shapedType = dyn_cast<ShapedType>(op.getType());
- if (shapedType && !shapedType.hasStaticShape())
- return failure();
-
ImplicitLocOpBuilder b(op->getLoc(), rewriter);
Value operand = op.getOperand();
Type opType = operand.getType();
+
+ if (isDynamicShaped(opType))
+ return failure();
+
Type operandETy = getElementTypeOrSelf(opType);
FloatType floatTy = llvm::dyn_cast<FloatType>(operandETy);
const llvm::fltSemantics &semantics = floatTy.getFloatSemantics();
@@ -449,6 +476,10 @@ static LogicalResult convertExp2fOp(math::Exp2Op op,
ImplicitLocOpBuilder b(op->getLoc(), rewriter);
Value operand = op.getOperand();
Type opType = operand.getType();
+
+ if (isDynamicShaped(opType))
+ return failure();
+
Value ln2 = createFloatConst(op->getLoc(), opType, llvm::numbers::ln2, b);
Value mult = arith::MulFOp::create(b, opType, operand, ln2);
Value exp = math::ExpOp::create(b, op->getLoc(), mult);
@@ -464,6 +495,9 @@ static LogicalResult convertRoundOp(math::RoundOp op,
Type opType = operand.getType();
Type opEType = getElementTypeOrSelf(opType);
+ if (isDynamicShaped(opType))
+ return failure();
+
if (!opEType.isF32()) {
return rewriter.notifyMatchFailure(op, "not a round of f32.");
}
@@ -523,6 +557,9 @@ static LogicalResult convertCtlzOp(math::CountLeadingZerosOp op,
auto eTy = getElementTypeOrSelf(operandTy);
Location loc = op.getLoc();
+ if (isDynamicShaped(operandTy))
+ return failure();
+
// Only expand for integer or float element types (index has no fixed bitwidth).
if (!eTy.isIntOrFloat()) {
return rewriter.notifyMatchFailure(op, "ctlz expansion only supports int or float types");
@@ -574,6 +611,9 @@ static LogicalResult convertRoundEvenOp(math::RoundEvenOp op,
Type operandETy = getElementTypeOrSelf(operandTy);
Type resultETy = getElementTypeOrSelf(resultTy);
+ if (isDynamicShaped(operandTy))
+ return failure();
+
if (!isa<FloatType>(operandETy) || !isa<FloatType>(resultETy)) {
return rewriter.notifyMatchFailure(op, "not a roundeven of f16 or f32.");
}
@@ -698,12 +738,10 @@ static LogicalResult convertRoundEvenOp(math::RoundEvenOp op,
// Convert `math.rsqrt` into `arith.divf` + `math.sqrt`
static LogicalResult convertRsqrtOp(math::RsqrtOp op,
PatternRewriter &rewriter) {
-
auto operand = op.getOperand();
auto operandTy = operand.getType();
- // Operand type must be shatic shaped type to create const float.
- auto shapedOperandType = dyn_cast<ShapedType>(operandTy);
- if (shapedOperandType && !shapedOperandType.hasStaticShape())
+
+ if (isDynamicShaped(operandTy))
return failure();
auto eTy = getElementTypeOrSelf(operandTy);
diff --git a/mlir/test/Dialect/Math/expand-math.mlir b/mlir/test/Dialect/Math/expand-math.mlir
index 126270ca40130..a483575793ab7 100644
--- a/mlir/test/Dialect/Math/expand-math.mlir
+++ b/mlir/test/Dialect/Math/expand-math.mlir
@@ -891,3 +891,263 @@ func.func @clampf_vector_op(%arg: vector<3x4xf32>, %min: vector<3x4xf32>, %max:
%a = math.clampf %arg to [%min, %max] fastmath<fast> : vector<3x4xf32>
return %a: vector<3x4xf32>
}
+
+// -----
+
+// CHECK-LABEL: func.func @non_static_shape_sinh_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
+// CHECK-SAME: -> tensor<?xf32>
+// CHECK: %[[OP:.*]] = math.sinh %[[ARG]] : tensor<?xf32>
+// CHECK: return %[[OP]] : tensor<?xf32>
+
+func.func @non_static_shape_sinh_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
+ %a = math.sinh %arg : tensor<?xf32>
+ return %a: tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unranked_sinh_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>)
+// CHECK-SAME: -> tensor<*xf32>
+// CHECK: %[[OP:.*]] = math.sinh %[[ARG]] : tensor<*xf32>
+// CHECK: return %[[OP]] : tensor<*xf32>
+
+func.func @unranked_sinh_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
+ %a = math.sinh %arg : tensor<*xf32>
+ return %a: tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @non_static_shape_cosh_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
+// CHECK-SAME: -> tensor<?xf32>
+// CHECK: %[[OP:.*]] = math.cosh %[[ARG]] : tensor<?xf32>
+// CHECK: return %[[OP]] : tensor<?xf32>
+
+func.func @non_static_shape_cosh_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
+ %a = math.cosh %arg : tensor<?xf32>
+ return %a: tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unranked_cosh_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>)
+// CHECK-SAME: -> tensor<*xf32>
+// CHECK: %[[OP:.*]] = math.cosh %[[ARG]] : tensor<*xf32>
+// CHECK: return %[[OP]] : tensor<*xf32>
+
+func.func @unranked_cosh_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
+ %a = math.cosh %arg : tensor<*xf32>
+ return %a: tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @non_static_shape_tanh_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
+// CHECK-SAME: -> tensor<?xf32>
+// CHECK: %[[OP:.*]] = math.tanh %[[ARG]] : tensor<?xf32>
+// CHECK: return %[[OP]] : tensor<?xf32>
+
+func.func @non_static_shape_tanh_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
+ %a = math.tanh %arg : tensor<?xf32>
+ return %a: tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unranked_tanh_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>)
+// CHECK-SAME: -> tensor<*xf32>
+// CHECK: %[[OP:.*]] = math.tanh %[[ARG]] : tensor<*xf32>
+// CHECK: return %[[OP]] : tensor<*xf32>
+
+func.func @unranked_tanh_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
+ %a = math.tanh %arg : tensor<*xf32>
+ return %a: tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @non_static_shape_asinh_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
+// CHECK-SAME: -> tensor<?xf32>
+// CHECK: %[[OP:.*]] = math.asinh %[[ARG]] : tensor<?xf32>
+// CHECK: return %[[OP]] : tensor<?xf32>
+
+func.func @non_static_shape_asinh_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
+ %a = math.asinh %arg : tensor<?xf32>
+ return %a: tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unranked_asinh_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>)
+// CHECK-SAME: -> tensor<*xf32>
+// CHECK: %[[OP:.*]] = math.asinh %[[ARG]] : tensor<*xf32>
+// CHECK: return %[[OP]] : tensor<*xf32>
+
+func.func @unranked_asinh_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
+ %a = math.asinh %arg : tensor<*xf32>
+ return %a: tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @non_static_shape_acosh_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
+// CHECK-SAME: -> tensor<?xf32>
+// CHECK: %[[OP:.*]] = math.acosh %[[ARG]] : tensor<?xf32>
+// CHECK: return %[[OP]] : tensor<?xf32>
+
+func.func @non_static_shape_acosh_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
+ %a = math.acosh %arg : tensor<?xf32>
+ return %a: tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unranked_acosh_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>)
+// CHECK-SAME: -> tensor<*xf32>
+// CHECK: %[[OP:.*]] = math.acosh %[[ARG]] : tensor<*xf32>
+// CHECK: return %[[OP]] : tensor<*xf32>
+
+func.func @unranked_acosh_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
+ %a = math.acosh %arg : tensor<*xf32>
+ return %a: tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @non_static_shape_atanh_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
+// CHECK-SAME: -> tensor<?xf32>
+// CHECK: %[[OP:.*]] = math.atanh %[[ARG]] : tensor<?xf32>
+// CHECK: return %[[OP]] : tensor<?xf32>
+
+func.func @non_static_shape_atanh_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
+ %a = math.atanh %arg : tensor<?xf32>
+ return %a: tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unranked_atanh_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>)
+// CHECK-SAME: -> tensor<*xf32>
+// CHECK: %[[OP:.*]] = math.atanh %[[ARG]] : tensor<*xf32>
+// CHECK: return %[[OP]] : tensor<*xf32>
+
+func.func @unranked_atanh_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
+ %a = math.atanh %arg : tensor<*xf32>
+ return %a: tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @non_static_shape_exp2_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
+// CHECK-SAME: -> tensor<?xf32>
+// CHECK: %[[OP:.*]] = math.exp2 %[[ARG]] : tensor<?xf32>
+// CHECK: return %[[OP]] : tensor<?xf32>
+
+func.func @non_static_shape_exp2_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
+ %a = math.exp2 %arg : tensor<?xf32>
+ return %a: tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unranked_exp2_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>)
+// CHECK-SAME: -> tensor<*xf32>
+// CHECK: %[[OP:.*]] = math.exp2 %[[ARG]] : tensor<*xf32>
+// CHECK: return %[[OP]] : tensor<*xf32>
+
+func.func @unranked_exp2_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
+ %a = math.exp2 %arg : tensor<*xf32>
+ return %a: tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @non_static_shape_round_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
+// CHECK-SAME: -> tensor<?xf32>
+// CHECK: %[[OP:.*]] = math.round %[[ARG]] : tensor<?xf32>
+// CHECK: return %[[OP]] : tensor<?xf32>
+
+func.func @non_static_shape_round_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
+ %a = math.round %arg : tensor<?xf32>
+ return %a: tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unranked_round_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>)
+// CHECK-SAME: -> tensor<*xf32>
+// CHECK: %[[OP:.*]] = math.round %[[ARG]] : tensor<*xf32>
+// CHECK: return %[[OP]] : tensor<*xf32>
+
+func.func @unranked_round_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
+ %a = math.round %arg : tensor<*xf32>
+ return %a: tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @non_static_shape_roundeven_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
+// CHECK-SAME: -> tensor<?xf32>
+// CHECK: %[[OP:.*]] = math.roundeven %[[ARG]] : tensor<?xf32>
+// CHECK: return %[[OP]] : tensor<?xf32>
+
+func.func @non_static_shape_roundeven_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
+ %a = math.roundeven %arg : tensor<?xf32>
+ return %a: tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unranked_roundeven_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>)
+// CHECK-SAME: -> tensor<*xf32>
+// CHECK: %[[OP:.*]] = math.roundeven %[[ARG]] : tensor<*xf32>
+// CHECK: return %[[OP]] : tensor<*xf32>
+
+func.func @unranked_roundeven_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
+ %a = math.roundeven %arg : tensor<*xf32>
+ return %a: tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @non_static_shape_ctlz_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<?xi32>)
+// CHECK-SAME: -> tensor<?xi32>
+// CHECK: %[[OP:.*]] = math.ctlz %[[ARG]] : tensor<?xi32>
+// CHECK: return %[[OP]] : tensor<?xi32>
+
+func.func @non_static_shape_ctlz_op(%arg: tensor<?xi32>) -> tensor<?xi32>{
+ %a = math.ctlz %arg : tensor<?xi32>
+ return %a: tensor<?xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unranked_ctlz_op
+// CHECK-SAME: (%[[ARG:.*]]: tensor<*xi32>)
+// CHECK-SAME: -> tensor<*xi32>
+// CHECK: %[[OP:.*]] = math.ctlz %[[ARG]] : tensor<*xi32>
+// CHECK: return %[[OP]] : tensor<*xi32>
+
+func.func @unranked_ctlz_op(%arg: tensor<*xi32>) -> tensor<*xi32>{
+ %a = math.ctlz %arg : tensor<*xi32>
+ return %a: tensor<*xi32>
+}
>From 6402001ba25de0b17ed77d8b9e7e400825b3ebec Mon Sep 17 00:00:00 2001
From: Peruere1828 <fmyh1828 at gmail.com>
Date: Fri, 26 Jun 2026 04:01:38 +0000
Subject: [PATCH 2/3] fix: handle dynamic-shaped tensors via scalar+splat in
createFloatConst
Modify createFloatConst to support dynamically-shaped ranked tensors
by creating a scalar constant and broadcasting via tensor.dim + tensor.splat
instead of asserting on DenseElementsAttr.
For ranked dynamic shapes (tensor<?xf32>), the expansion now works correctly.
Unranked tensors (tensor<*xf32>) still pass through, since tensor.splat
requires ranked tensor types.
Simple ops (sinh, cosh, tanh, asinh, acosh, atanh, exp2, rsqrt) are properly
expand on dynamic-shaped tensors instead of bailing out.
Complex ops (ceil, round, roundeven, ctlz) that also use createIntConst
retain the full isDynamicShaped bailout for now.
---
.../mlir/Dialect/Math/Transforms/Passes.td | 2 +-
.../Dialect/Math/Transforms/CMakeLists.txt | 1 +
.../lib/Dialect/Math/Transforms/ExpandOps.cpp | 79 ++++++++++++-------
mlir/test/Dialect/Math/expand-math.mlir | 61 ++++++++++----
4 files changed, 98 insertions(+), 45 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Math/Transforms/Passes.td b/mlir/include/mlir/Dialect/Math/Transforms/Passes.td
index 48346abd84285..4e7c0330e7287 100644
--- a/mlir/include/mlir/Dialect/Math/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Math/Transforms/Passes.td
@@ -57,7 +57,7 @@ def MathExpandOpsPass : Pass<"math-expand-ops"> {
`math.acosh` operations. If the list is empty, then all expansions are
applied.
}];
- let dependentDialects = ["arith::ArithDialect"];
+ let dependentDialects = ["arith::ArithDialect", "mlir::tensor::TensorDialect"];
let options = [
ListOption<"opMnemonics", "ops", "std::string",
"Operations to expand.">
diff --git a/mlir/lib/Dialect/Math/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Math/Transforms/CMakeLists.txt
index dc7e78288eb29..3cbfe03d784ab 100644
--- a/mlir/lib/Dialect/Math/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Math/Transforms/CMakeLists.txt
@@ -20,6 +20,7 @@ add_mlir_dialect_library(MLIRMathTransforms
MLIRMathDialect
MLIRSCFDialect
MLIRPass
+ MLIRTensorDialect
MLIRTransforms
MLIRX86Dialect
MLIRVectorDialect
diff --git a/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
index 27bcaee5c8303..7b99515ea99bc 100644
--- a/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
@@ -13,6 +13,7 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/Math/Transforms/Passes.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/TypeUtilities.h"
@@ -34,9 +35,18 @@ static bool isDynamicShaped(Type type) {
return shapedTy && !shapedTy.hasStaticShape();
}
-/// Create a float constant.
+/// Returns true if the type is an unranked shaped type (e.g., tensor<*xf32>).
+/// Unranked types can't be expanded because tensor.splat requires ranked types.
+static bool isUnrankedShaped(Type type) {
+ auto shapedTy = dyn_cast<ShapedType>(type);
+ return shapedTy && !shapedTy.hasRank();
+}
+
+/// Create a float constant. For dynamically-shaped tensors, creates a scalar
+/// constant and uses tensor.dim + tensor.splat to broadcast to the target shape.
+/// The optional `dynamicShapeRef` provides the runtime dimension sizes.
static Value createFloatConst(Location loc, Type type, APFloat value,
- OpBuilder &b) {
+ OpBuilder &b, Value dynamicShapeRef = Value()) {
bool losesInfo = false;
auto eltType = getElementTypeOrSelf(type);
// Convert double to the given `FloatType` with round-to-nearest-ties-to-even.
@@ -44,16 +54,27 @@ static Value createFloatConst(Location loc, Type type, APFloat value,
APFloat::rmNearestTiesToEven, &losesInfo);
auto attr = b.getFloatAttr(eltType, value);
if (auto shapedTy = dyn_cast<ShapedType>(type)) {
- return arith::ConstantOp::create(b, loc,
- DenseElementsAttr::get(shapedTy, attr));
+ if (shapedTy.hasStaticShape())
+ return arith::ConstantOp::create(b, loc,
+ DenseElementsAttr::get(shapedTy, attr));
+
+ // Dynamic shape: create scalar constant and splat to the target shape.
+ Value scalar = arith::ConstantOp::create(b, loc, eltType, attr);
+ SmallVector<Value> dynamicSizes;
+ for (int64_t i = 0; i < shapedTy.getRank(); ++i) {
+ if (shapedTy.isDynamicDim(i))
+ dynamicSizes.push_back(
+ tensor::DimOp::create(b, loc, dynamicShapeRef, i));
+ }
+ return tensor::SplatOp::create(b, loc, type, scalar, dynamicSizes);
}
return arith::ConstantOp::create(b, loc, attr);
}
static Value createFloatConst(Location loc, Type type, double value,
- OpBuilder &b) {
- return createFloatConst(loc, type, APFloat(value), b);
+ OpBuilder &b, Value dynamicShapeRef = Value()) {
+ return createFloatConst(loc, type, APFloat(value), b, dynamicShapeRef);
}
/// Create an integer constant.
@@ -86,14 +107,14 @@ static LogicalResult convertSinhOp(math::SinhOp op, PatternRewriter &rewriter) {
Value operand = op.getOperand();
Type opType = operand.getType();
- if (isDynamicShaped(opType))
+ if (isUnrankedShaped(opType))
return failure();
Value exp = math::ExpOp::create(b, operand);
Value neg = arith::NegFOp::create(b, operand);
Value nexp = math::ExpOp::create(b, neg);
Value sub = arith::SubFOp::create(b, exp, nexp);
- Value half = createFloatConst(op->getLoc(), opType, 0.5, rewriter);
+ Value half = createFloatConst(op->getLoc(), opType, 0.5, rewriter, operand);
Value res = arith::MulFOp::create(b, sub, half);
rewriter.replaceOp(op, res);
return success();
@@ -105,14 +126,14 @@ static LogicalResult convertCoshOp(math::CoshOp op, PatternRewriter &rewriter) {
Value operand = op.getOperand();
Type opType = operand.getType();
- if (isDynamicShaped(opType))
+ if (isUnrankedShaped(opType))
return failure();
Value exp = math::ExpOp::create(b, operand);
Value neg = arith::NegFOp::create(b, operand);
Value nexp = math::ExpOp::create(b, neg);
Value add = arith::AddFOp::create(b, exp, nexp);
- Value half = createFloatConst(op->getLoc(), opType, 0.5, rewriter);
+ Value half = createFloatConst(op->getLoc(), opType, 0.5, rewriter, operand);
Value res = arith::MulFOp::create(b, add, half);
rewriter.replaceOp(op, res);
return success();
@@ -127,19 +148,20 @@ static LogicalResult convertCoshOp(math::CoshOp op, PatternRewriter &rewriter) {
/// 1]. Expand the computation on the input `x * sign(x)`, then multiply the
/// result by `sign(x)` to retain sign of the real result.
static LogicalResult convertTanhOp(math::TanhOp op, PatternRewriter &rewriter) {
- auto floatType = op.getOperand().getType();
+ Value operand = op.getOperand();
+ auto floatType = operand.getType();
- if (isDynamicShaped(floatType))
+ if (isUnrankedShaped(floatType))
return failure();
Location loc = op.getLoc();
- Value zero = createFloatConst(loc, floatType, 0.0, rewriter);
- Value one = createFloatConst(loc, floatType, 1.0, rewriter);
- Value negTwo = createFloatConst(loc, floatType, -2.0, rewriter);
+ Value zero = createFloatConst(loc, floatType, 0.0, rewriter, operand);
+ Value one = createFloatConst(loc, floatType, 1.0, rewriter, operand);
+ Value negTwo = createFloatConst(loc, floatType, -2.0, rewriter, operand);
// Compute sign(x) = cast<float_type>(x < 0) * (-2) + 1
Value isNegative = arith::CmpFOp::create(
- rewriter, loc, arith::CmpFPredicate::OLT, op.getOperand(), zero);
+ rewriter, loc, arith::CmpFPredicate::OLT, operand, zero);
Value isNegativeFloat =
arith::UIToFPOp::create(rewriter, loc, floatType, isNegative);
Value isNegativeTimesNegTwo =
@@ -147,7 +169,7 @@ static LogicalResult convertTanhOp(math::TanhOp op, PatternRewriter &rewriter) {
Value sign = arith::AddFOp::create(rewriter, loc, isNegativeTimesNegTwo, one);
// Normalize input to positive value: y = sign(x) * x
- Value positiveX = arith::MulFOp::create(rewriter, loc, sign, op.getOperand());
+ Value positiveX = arith::MulFOp::create(rewriter, loc, sign, operand);
// Decompose on normalized input
Value negDoubledX = arith::MulFOp::create(rewriter, loc, negTwo, positiveX);
@@ -181,10 +203,10 @@ static LogicalResult convertAsinhOp(math::AsinhOp op,
Value operand = op.getOperand();
Type opType = operand.getType();
- if (isDynamicShaped(opType))
+ if (isUnrankedShaped(opType))
return failure();
- Value one = createFloatConst(op->getLoc(), opType, 1.0, rewriter);
+ Value one = createFloatConst(op->getLoc(), opType, 1.0, rewriter, operand);
Value fma = math::FmaOp::create(b, operand, operand, one);
Value sqrt = math::SqrtOp::create(b, fma);
Value add = arith::AddFOp::create(b, operand, sqrt);
@@ -200,10 +222,10 @@ static LogicalResult convertAcoshOp(math::AcoshOp op,
Value operand = op.getOperand();
Type opType = operand.getType();
- if (isDynamicShaped(opType))
+ if (isUnrankedShaped(opType))
return failure();
- Value negOne = createFloatConst(op->getLoc(), opType, -1.0, rewriter);
+ Value negOne = createFloatConst(op->getLoc(), opType, -1.0, rewriter, operand);
Value fma = math::FmaOp::create(b, operand, operand, negOne);
Value sqrt = math::SqrtOp::create(b, fma);
Value add = arith::AddFOp::create(b, operand, sqrt);
@@ -219,16 +241,16 @@ static LogicalResult convertAtanhOp(math::AtanhOp op,
Value operand = op.getOperand();
Type opType = operand.getType();
- if (isDynamicShaped(opType))
+ if (isUnrankedShaped(opType))
return failure();
- Value one = createFloatConst(op->getLoc(), opType, 1.0, rewriter);
+ Value one = createFloatConst(op->getLoc(), opType, 1.0, rewriter, operand);
Value add = arith::AddFOp::create(b, operand, one);
Value neg = arith::NegFOp::create(b, operand);
Value sub = arith::AddFOp::create(b, neg, one);
Value div = arith::DivFOp::create(b, add, sub);
Value log = math::LogOp::create(b, div);
- Value half = createFloatConst(op->getLoc(), opType, 0.5, rewriter);
+ Value half = createFloatConst(op->getLoc(), opType, 0.5, rewriter, operand);
Value res = arith::MulFOp::create(b, log, half);
rewriter.replaceOp(op, res);
return success();
@@ -477,10 +499,11 @@ static LogicalResult convertExp2fOp(math::Exp2Op op,
Value operand = op.getOperand();
Type opType = operand.getType();
- if (isDynamicShaped(opType))
+ if (isUnrankedShaped(opType))
return failure();
- Value ln2 = createFloatConst(op->getLoc(), opType, llvm::numbers::ln2, b);
+ Value ln2 = createFloatConst(op->getLoc(), opType, llvm::numbers::ln2, b,
+ operand);
Value mult = arith::MulFOp::create(b, opType, operand, ln2);
Value exp = math::ExpOp::create(b, op->getLoc(), mult);
rewriter.replaceOp(op, exp);
@@ -741,7 +764,7 @@ static LogicalResult convertRsqrtOp(math::RsqrtOp op,
auto operand = op.getOperand();
auto operandTy = operand.getType();
- if (isDynamicShaped(operandTy))
+ if (isUnrankedShaped(operandTy))
return failure();
auto eTy = getElementTypeOrSelf(operandTy);
@@ -749,7 +772,7 @@ static LogicalResult convertRsqrtOp(math::RsqrtOp op,
return failure();
Location loc = op->getLoc();
- auto constOneFloat = createFloatConst(loc, operandTy, 1.0, rewriter);
+ auto constOneFloat = createFloatConst(loc, operandTy, 1.0, rewriter, operand);
auto sqrtOp = math::SqrtOp::create(rewriter, loc, operand);
rewriter.replaceOpWithNewOp<arith::DivFOp>(op, constOneFloat, sqrtOp);
return success();
diff --git a/mlir/test/Dialect/Math/expand-math.mlir b/mlir/test/Dialect/Math/expand-math.mlir
index a483575793ab7..1abdd6a302276 100644
--- a/mlir/test/Dialect/Math/expand-math.mlir
+++ b/mlir/test/Dialect/Math/expand-math.mlir
@@ -847,8 +847,12 @@ func.func @unranked_ceil_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
// CHECK-LABEL: func.func @non_static_shape_rsqrt_op
// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
// CHECK-SAME: -> tensor<?xf32>
-// CHECK: %[[RSQRT:.*]] = math.rsqrt %[[ARG]] : tensor<?xf32>
-// CHECK: return %[[RSQRT]] : tensor<?xf32>
+// CHECK-DAG: arith.constant 1.000000e+00 : f32
+// CHECK-DAG: tensor.dim %[[ARG]]
+// CHECK-DAG: tensor.splat %{{.*}}[%{{.*}}] : tensor<?xf32>
+// CHECK-DAG: math.sqrt %[[ARG]] : tensor<?xf32>
+// CHECK: arith.divf %{{.*}}, %{{.*}} : tensor<?xf32>
+// CHECK: return %{{.*}} : tensor<?xf32>
func.func @non_static_shape_rsqrt_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
%a = math.rsqrt %arg : tensor<?xf32>
@@ -897,8 +901,12 @@ func.func @clampf_vector_op(%arg: vector<3x4xf32>, %min: vector<3x4xf32>, %max:
// CHECK-LABEL: func.func @non_static_shape_sinh_op
// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
// CHECK-SAME: -> tensor<?xf32>
-// CHECK: %[[OP:.*]] = math.sinh %[[ARG]] : tensor<?xf32>
-// CHECK: return %[[OP]] : tensor<?xf32>
+// CHECK-DAG: math.exp %[[ARG]] : tensor<?xf32>
+// CHECK-DAG: arith.negf %[[ARG]] : tensor<?xf32>
+// CHECK-DAG: arith.constant 5.000000e-01 : f32
+// CHECK-DAG: tensor.dim %[[ARG]]
+// CHECK: tensor.splat %{{.*}}[%{{.*}}] : tensor<?xf32>
+// CHECK: return %{{.*}} : tensor<?xf32>
func.func @non_static_shape_sinh_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
%a = math.sinh %arg : tensor<?xf32>
@@ -923,8 +931,12 @@ func.func @unranked_sinh_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
// CHECK-LABEL: func.func @non_static_shape_cosh_op
// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
// CHECK-SAME: -> tensor<?xf32>
-// CHECK: %[[OP:.*]] = math.cosh %[[ARG]] : tensor<?xf32>
-// CHECK: return %[[OP]] : tensor<?xf32>
+// CHECK-DAG: math.exp %[[ARG]] : tensor<?xf32>
+// CHECK-DAG: arith.negf %[[ARG]] : tensor<?xf32>
+// CHECK-DAG: arith.constant 5.000000e-01 : f32
+// CHECK-DAG: tensor.dim %[[ARG]]
+// CHECK: tensor.splat %{{.*}}[%{{.*}}] : tensor<?xf32>
+// CHECK: return %{{.*}} : tensor<?xf32>
func.func @non_static_shape_cosh_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
%a = math.cosh %arg : tensor<?xf32>
@@ -949,8 +961,12 @@ func.func @unranked_cosh_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
// CHECK-LABEL: func.func @non_static_shape_tanh_op
// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
// CHECK-SAME: -> tensor<?xf32>
-// CHECK: %[[OP:.*]] = math.tanh %[[ARG]] : tensor<?xf32>
-// CHECK: return %[[OP]] : tensor<?xf32>
+// CHECK-DAG: arith.constant 0.000000e+00 : f32
+// CHECK-DAG: arith.constant 1.000000e+00 : f32
+// CHECK-DAG: arith.constant -2.000000e+00 : f32
+// CHECK-DAG: tensor.dim %[[ARG]]
+// CHECK: tensor.splat %{{.*}}[%{{.*}}] : tensor<?xf32>
+// CHECK: return %{{.*}} : tensor<?xf32>
func.func @non_static_shape_tanh_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
%a = math.tanh %arg : tensor<?xf32>
@@ -975,8 +991,11 @@ func.func @unranked_tanh_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
// CHECK-LABEL: func.func @non_static_shape_asinh_op
// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
// CHECK-SAME: -> tensor<?xf32>
-// CHECK: %[[OP:.*]] = math.asinh %[[ARG]] : tensor<?xf32>
-// CHECK: return %[[OP]] : tensor<?xf32>
+// CHECK-DAG: arith.constant 1.000000e+00 : f32
+// CHECK-DAG: tensor.dim %[[ARG]]
+// CHECK: tensor.splat %{{.*}}[%{{.*}}] : tensor<?xf32>
+// CHECK: math.log %{{.*}} : tensor<?xf32>
+// CHECK: return %{{.*}} : tensor<?xf32>
func.func @non_static_shape_asinh_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
%a = math.asinh %arg : tensor<?xf32>
@@ -1001,8 +1020,11 @@ func.func @unranked_asinh_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
// CHECK-LABEL: func.func @non_static_shape_acosh_op
// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
// CHECK-SAME: -> tensor<?xf32>
-// CHECK: %[[OP:.*]] = math.acosh %[[ARG]] : tensor<?xf32>
-// CHECK: return %[[OP]] : tensor<?xf32>
+// CHECK-DAG: arith.constant -1.000000e+00 : f32
+// CHECK-DAG: tensor.dim %[[ARG]]
+// CHECK: tensor.splat %{{.*}}[%{{.*}}] : tensor<?xf32>
+// CHECK: math.log %{{.*}} : tensor<?xf32>
+// CHECK: return %{{.*}} : tensor<?xf32>
func.func @non_static_shape_acosh_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
%a = math.acosh %arg : tensor<?xf32>
@@ -1027,8 +1049,12 @@ func.func @unranked_acosh_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
// CHECK-LABEL: func.func @non_static_shape_atanh_op
// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
// CHECK-SAME: -> tensor<?xf32>
-// CHECK: %[[OP:.*]] = math.atanh %[[ARG]] : tensor<?xf32>
-// CHECK: return %[[OP]] : tensor<?xf32>
+// CHECK-DAG: arith.constant 1.000000e+00 : f32
+// CHECK-DAG: arith.constant 5.000000e-01 : f32
+// CHECK-DAG: tensor.dim %[[ARG]]
+// CHECK: tensor.splat %{{.*}}[%{{.*}}] : tensor<?xf32>
+// CHECK: math.log %{{.*}} : tensor<?xf32>
+// CHECK: return %{{.*}} : tensor<?xf32>
func.func @non_static_shape_atanh_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
%a = math.atanh %arg : tensor<?xf32>
@@ -1053,8 +1079,11 @@ func.func @unranked_atanh_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
// CHECK-LABEL: func.func @non_static_shape_exp2_op
// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
// CHECK-SAME: -> tensor<?xf32>
-// CHECK: %[[OP:.*]] = math.exp2 %[[ARG]] : tensor<?xf32>
-// CHECK: return %[[OP]] : tensor<?xf32>
+// CHECK-DAG: arith.constant 0.6931471{{.*}} : f32
+// CHECK-DAG: tensor.dim %[[ARG]]
+// CHECK: tensor.splat %{{.*}}[%{{.*}}] : tensor<?xf32>
+// CHECK: math.exp %{{.*}} : tensor<?xf32>
+// CHECK: return %{{.*}} : tensor<?xf32>
func.func @non_static_shape_exp2_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
%a = math.exp2 %arg : tensor<?xf32>
>From 382088f216fbdaba27bff9b2343f91e523bebe56 Mon Sep 17 00:00:00 2001
From: Peruere1828 <fmyh1828 at gmail.com>
Date: Fri, 26 Jun 2026 08:48:16 +0000
Subject: [PATCH 3/3] fix: handle dynamic-shaped tensors via scalar+splat in
createIntConst
Modify createIntConst to support dynamically-shaped ranked tensors
by creating a scalar constant and broadcasting via tensor.dim + tensor.splat
instead of asserting on DenseElementsAttr. Same as createFloatConst.
For ranked dynamic shapes (tensor<?xf32>), the expansion now works correctly.
Unranked tensors (tensor<*xf32>) still pass through, since tensor.splat
requires ranked tensor types.
Now all ops (sinh, cosh, tanh, asinh, acosh, atanh, exp2, rsqrt, ceil,
round, roundeven, ctlz) are properly expand on dynamic-shaped tensors
instead of bailing out.
---
.../lib/Dialect/Math/Transforms/ExpandOps.cpp | 96 +++++++++++--------
mlir/test/Dialect/Math/expand-math.mlir | 34 +++++--
2 files changed, 80 insertions(+), 50 deletions(-)
diff --git a/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
index 7b99515ea99bc..15188208a8b4c 100644
--- a/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
@@ -26,15 +26,6 @@ namespace mlir::math {
#include "mlir/Dialect/Math/Transforms/Passes.h.inc"
} // namespace mlir::math
-/// Check if the type is a shaped type with dynamic shape
-/// (e.g., tensor<?xf32>). Returns false for scalar types and static-shaped
-/// types, allowing expansion to proceed. Only dynamic-shaped tensors/vectors
-/// are rejected since the generated constant folding assumes static shapes.
-static bool isDynamicShaped(Type type) {
- auto shapedTy = dyn_cast<ShapedType>(type);
- return shapedTy && !shapedTy.hasStaticShape();
-}
-
/// Returns true if the type is an unranked shaped type (e.g., tensor<*xf32>).
/// Unranked types can't be expanded because tensor.splat requires ranked types.
static bool isUnrankedShaped(Type type) {
@@ -77,13 +68,27 @@ static Value createFloatConst(Location loc, Type type, double value,
return createFloatConst(loc, type, APFloat(value), b, dynamicShapeRef);
}
-/// Create an integer constant.
+/// Create an integer constant. For dynamically-shaped tensors, creates a scalar
+/// constant and uses tensor.dim + tensor.splat to broadcast to the target shape.
+/// The optional `dynamicShapeRef` provides the runtime dimension sizes.
static Value createIntConst(Location loc, Type type, int64_t value,
- OpBuilder &b) {
- auto attr = b.getIntegerAttr(getElementTypeOrSelf(type), value);
+ OpBuilder &b, Value dynamicShapeRef = Value()) {
+ auto eltType = getElementTypeOrSelf(type);
+ auto attr = b.getIntegerAttr(eltType, value);
if (auto shapedTy = dyn_cast<ShapedType>(type)) {
- return arith::ConstantOp::create(b, loc,
- DenseElementsAttr::get(shapedTy, attr));
+ if (shapedTy.hasStaticShape())
+ return arith::ConstantOp::create(b, loc,
+ DenseElementsAttr::get(shapedTy, attr));
+
+ // Dynamic shape: create scalar constant and splat to the target shape.
+ Value scalar = arith::ConstantOp::create(b, loc, eltType, attr);
+ SmallVector<Value> dynamicSizes;
+ for (int64_t i = 0; i < shapedTy.getRank(); ++i) {
+ if (shapedTy.isDynamicDim(i))
+ dynamicSizes.push_back(
+ tensor::DimOp::create(b, loc, dynamicShapeRef, i));
+ }
+ return tensor::SplatOp::create(b, loc, type, scalar, dynamicSizes);
}
return arith::ConstantOp::create(b, loc, attr);
@@ -278,7 +283,7 @@ static LogicalResult convertCeilOp(math::CeilOp op, PatternRewriter &rewriter) {
Value operand = op.getOperand();
Type opType = operand.getType();
- if (isDynamicShaped(opType))
+ if (isUnrankedShaped(opType))
return failure();
Type operandETy = getElementTypeOrSelf(opType);
@@ -306,12 +311,13 @@ static LogicalResult convertCeilOp(math::CeilOp op, PatternRewriter &rewriter) {
// For all such cases, `ceilf(x)` is defined to return `x` directly.
Value operandBitcast = arith::BitcastOp::create(b, iTy, operand);
Value cMask = createIntConst(
- op->getLoc(), iTy, static_cast<int64_t>((1ull << (bitWidth - 1)) - 1), b);
+ op->getLoc(), iTy, static_cast<int64_t>((1ull << (bitWidth - 1)) - 1), b,
+ operand);
Value unsignedBits = arith::AndIOp::create(b, operandBitcast, cMask);
Value cThreshold = createIntConst(
op->getLoc(), iTy,
static_cast<int64_t>((uint64_t(bias + mantissaWidth)) << mantissaWidth),
- b);
+ b, operand);
Value isLargeExp = arith::CmpIOp::create(b, arith::CmpIPredicate::uge,
unsignedBits, cThreshold);
Value isSpecialValOrLargeVal = isLargeExp;
@@ -320,7 +326,8 @@ static LogicalResult convertCeilOp(math::CeilOp op, PatternRewriter &rewriter) {
// all 0s in the exponent and mantissa, therefore requires an explicit check.
if (hasNegativeZeroNaNEncoding) {
Value cNegZeroBits = createIntConst(
- op->getLoc(), iTy, static_cast<int64_t>(1ull << (bitWidth - 1)), b);
+ op->getLoc(), iTy, static_cast<int64_t>(1ull << (bitWidth - 1)), b,
+ operand);
Value isNegZeroEncoding = arith::CmpIOp::create(
b, arith::CmpIPredicate::eq, operandBitcast, cNegZeroBits);
isSpecialValOrLargeVal =
@@ -330,8 +337,8 @@ static LogicalResult convertCeilOp(math::CeilOp op, PatternRewriter &rewriter) {
Value fpFixedConvert = createTruncatedFPValue(operand, b);
// Creating constants for later use.
- Value zero = createFloatConst(op->getLoc(), opType, 0.00, rewriter);
- Value one = createFloatConst(op->getLoc(), opType, 1.00, rewriter);
+ Value zero = createFloatConst(op->getLoc(), opType, 0.00, rewriter, operand);
+ Value one = createFloatConst(op->getLoc(), opType, 1.00, rewriter, operand);
Value gtCheck = arith::CmpFOp::create(b, arith::CmpFPredicate::OGT, operand,
fpFixedConvert);
@@ -518,7 +525,7 @@ static LogicalResult convertRoundOp(math::RoundOp op,
Type opType = operand.getType();
Type opEType = getElementTypeOrSelf(opType);
- if (isDynamicShaped(opType))
+ if (isUnrankedShaped(opType))
return failure();
if (!opEType.isF32()) {
@@ -529,10 +536,10 @@ static LogicalResult convertRoundOp(math::RoundOp op,
if (auto shapedTy = dyn_cast<ShapedType>(opType))
i32Ty = shapedTy.clone(i32Ty);
- Value half = createFloatConst(loc, opType, 0.5, b);
- Value c23 = createIntConst(loc, i32Ty, 23, b);
- Value c127 = createIntConst(loc, i32Ty, 127, b);
- Value expMask = createIntConst(loc, i32Ty, (1 << 8) - 1, b);
+ Value half = createFloatConst(loc, opType, 0.5, b, operand);
+ Value c23 = createIntConst(loc, i32Ty, 23, b, operand);
+ Value c127 = createIntConst(loc, i32Ty, 127, b, operand);
+ Value expMask = createIntConst(loc, i32Ty, (1 << 8) - 1, b, operand);
Value incrValue = math::CopySignOp::create(b, half, operand);
Value add = arith::AddFOp::create(b, opType, operand, incrValue);
@@ -580,7 +587,7 @@ static LogicalResult convertCtlzOp(math::CountLeadingZerosOp op,
auto eTy = getElementTypeOrSelf(operandTy);
Location loc = op.getLoc();
- if (isDynamicShaped(operandTy))
+ if (isUnrankedShaped(operandTy))
return failure();
// Only expand for integer or float element types (index has no fixed bitwidth).
@@ -598,11 +605,12 @@ static LogicalResult convertCtlzOp(math::CountLeadingZerosOp op,
}
Value x = operand;
- Value count = createIntConst(loc, operandTy, 0, rewriter);
+ Value count = createIntConst(loc, operandTy, 0, rewriter, operand);
for (int32_t bw = bitwidth; bw > 1; bw = bw / 2) {
auto half = bw / 2;
- auto bits = createIntConst(loc, operandTy, half, rewriter);
- auto mask = createIntConst(loc, operandTy, allbits >> half, rewriter);
+ auto bits = createIntConst(loc, operandTy, half, rewriter, operand);
+ auto mask = createIntConst(loc, operandTy, allbits >> half, rewriter,
+ operand);
Value pred = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::ule,
x, mask);
@@ -613,11 +621,11 @@ static LogicalResult convertCtlzOp(math::CountLeadingZerosOp op,
count = arith::SelectOp::create(rewriter, loc, pred, add, count);
}
- Value zero = createIntConst(loc, operandTy, 0, rewriter);
+ Value zero = createIntConst(loc, operandTy, 0, rewriter, operand);
Value pred = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq,
operand, zero);
- Value bwval = createIntConst(loc, operandTy, bitwidth, rewriter);
+ Value bwval = createIntConst(loc, operandTy, bitwidth, rewriter, operand);
Value sel = arith::SelectOp::create(rewriter, loc, pred, bwval, count);
rewriter.replaceOp(op, sel);
return success();
@@ -634,7 +642,7 @@ static LogicalResult convertRoundEvenOp(math::RoundEvenOp op,
Type operandETy = getElementTypeOrSelf(operandTy);
Type resultETy = getElementTypeOrSelf(resultTy);
- if (isDynamicShaped(operandTy))
+ if (isUnrankedShaped(operandTy))
return failure();
if (!isa<FloatType>(operandETy) || !isa<FloatType>(resultETy)) {
@@ -657,16 +665,20 @@ static LogicalResult convertRoundEvenOp(math::RoundEvenOp op,
// f64: 1 bit sign | 11 bits exponent | 52 bits mantissa.
// f32: 1 bit sign | 8 bits exponent | 23 bits mantissa.
// f16: 1 bit sign | 5 bits exponent | 10 bits mantissa.
- Value c1Float = createFloatConst(loc, fTy, 1.0, b);
- Value c0 = createIntConst(loc, iTy, 0, b);
- Value c1 = createIntConst(loc, iTy, 1, b);
- Value cNeg1 = createIntConst(loc, iTy, -1, b);
- Value c23 = createIntConst(loc, iTy, mantissaWidth, b);
- Value c31 = createIntConst(loc, iTy, bitWidth - 1, b);
- Value c127 = createIntConst(loc, iTy, (1ull << (exponentWidth - 1)) - 1, b);
- Value c2To22 = createIntConst(loc, iTy, 1ull << (mantissaWidth - 1), b);
- Value c23Mask = createIntConst(loc, iTy, (1ull << mantissaWidth) - 1, b);
- Value expMask = createIntConst(loc, iTy, (1ull << exponentWidth) - 1, b);
+ Value c1Float = createFloatConst(loc, fTy, 1.0, b, operand);
+ Value c0 = createIntConst(loc, iTy, 0, b, operand);
+ Value c1 = createIntConst(loc, iTy, 1, b, operand);
+ Value cNeg1 = createIntConst(loc, iTy, -1, b, operand);
+ Value c23 = createIntConst(loc, iTy, mantissaWidth, b, operand);
+ Value c31 = createIntConst(loc, iTy, bitWidth - 1, b, operand);
+ Value c127 = createIntConst(loc, iTy, (1ull << (exponentWidth - 1)) - 1, b,
+ operand);
+ Value c2To22 = createIntConst(loc, iTy, 1ull << (mantissaWidth - 1), b,
+ operand);
+ Value c23Mask = createIntConst(loc, iTy, (1ull << mantissaWidth) - 1, b,
+ operand);
+ Value expMask = createIntConst(loc, iTy, (1ull << exponentWidth) - 1, b,
+ operand);
Value operandBitcast = arith::BitcastOp::create(b, iTy, operand);
Value round = math::RoundOp::create(b, operand);
diff --git a/mlir/test/Dialect/Math/expand-math.mlir b/mlir/test/Dialect/Math/expand-math.mlir
index 1abdd6a302276..84b0bb0aa314a 100644
--- a/mlir/test/Dialect/Math/expand-math.mlir
+++ b/mlir/test/Dialect/Math/expand-math.mlir
@@ -821,8 +821,12 @@ func.func @rsqrt_tns(%float: tensor<5x8xf32>) -> (tensor<5x8xf32>) {
// CHECK-LABEL: func.func @non_static_shape_ceil_op
// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
// CHECK-SAME: -> tensor<?xf32>
-// CHECK: %[[CEIL:.*]] = math.ceil %[[ARG]] : tensor<?xf32>
-// CHECK: return %[[CEIL]] : tensor<?xf32>
+// CHECK-DAG: arith.constant 0.000000e+00 : f32
+// CHECK-DAG: arith.constant 1.000000e+00 : f32
+// CHECK-DAG: tensor.dim %[[ARG]]
+// CHECK: tensor.splat %{{.*}}[%{{.*}}] : tensor<?xf32>
+// CHECK-NOT: math.ceil
+// CHECK: return %{{.*}} : tensor<?xf32>
func.func @non_static_shape_ceil_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
%a = math.ceil %arg : tensor<?xf32>
@@ -1108,8 +1112,14 @@ func.func @unranked_exp2_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
// CHECK-LABEL: func.func @non_static_shape_round_op
// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
// CHECK-SAME: -> tensor<?xf32>
-// CHECK: %[[OP:.*]] = math.round %[[ARG]] : tensor<?xf32>
-// CHECK: return %[[OP]] : tensor<?xf32>
+// CHECK-DAG: arith.constant 5.000000e-01 : f32
+// CHECK-DAG: arith.constant 23 : i32
+// CHECK-DAG: arith.constant 127 : i32
+// CHECK-DAG: arith.constant 255 : i32
+// CHECK-DAG: tensor.dim %[[ARG]]
+// CHECK: tensor.splat %{{.*}}[%{{.*}}] : tensor<?xf32>
+// CHECK-NOT: math.round
+// CHECK: return %{{.*}} : tensor<?xf32>
func.func @non_static_shape_round_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
%a = math.round %arg : tensor<?xf32>
@@ -1134,8 +1144,12 @@ func.func @unranked_round_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
// CHECK-LABEL: func.func @non_static_shape_roundeven_op
// CHECK-SAME: (%[[ARG:.*]]: tensor<?xf32>)
// CHECK-SAME: -> tensor<?xf32>
-// CHECK: %[[OP:.*]] = math.roundeven %[[ARG]] : tensor<?xf32>
-// CHECK: return %[[OP]] : tensor<?xf32>
+// CHECK-DAG: arith.constant 1.000000e+00 : f32
+// CHECK-DAG: arith.constant 0 : i32
+// CHECK-DAG: tensor.dim %[[ARG]]
+// CHECK: tensor.splat %{{.*}}[%{{.*}}] : tensor<?xf32>
+// CHECK-NOT: math.roundeven
+// CHECK: return %{{.*}} : tensor<?xf32>
func.func @non_static_shape_roundeven_op(%arg: tensor<?xf32>) -> tensor<?xf32>{
%a = math.roundeven %arg : tensor<?xf32>
@@ -1160,8 +1174,12 @@ func.func @unranked_roundeven_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
// CHECK-LABEL: func.func @non_static_shape_ctlz_op
// CHECK-SAME: (%[[ARG:.*]]: tensor<?xi32>)
// CHECK-SAME: -> tensor<?xi32>
-// CHECK: %[[OP:.*]] = math.ctlz %[[ARG]] : tensor<?xi32>
-// CHECK: return %[[OP]] : tensor<?xi32>
+// CHECK-DAG: arith.constant 0 : i32
+// CHECK-DAG: arith.constant 16 : i32
+// CHECK-DAG: tensor.dim %[[ARG]]
+// CHECK: tensor.splat %{{.*}}[%{{.*}}] : tensor<?xi32>
+// CHECK-NOT: math.ctlz
+// CHECK: return %{{.*}} : tensor<?xi32>
func.func @non_static_shape_ctlz_op(%arg: tensor<?xi32>) -> tensor<?xi32>{
%a = math.ctlz %arg : tensor<?xi32>
More information about the Mlir-commits
mailing list