[Mlir-commits] [mlir] [MLIR][Math] Handle dynamic-shaped tensors without assertion (PR #204510)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jun 17 22:38:14 PDT 2026
https://github.com/Peruere1828 created https://github.com/llvm/llvm-project/pull/204510
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
>From a2a22711ea7930405af11bd185480cacdde9c320 Mon Sep 17 00:00:00 2001
From: Peruere1828 <fmyh1828 at gmail.com>
Date: Thu, 18 Jun 2026 05:27:10 +0000
Subject: [PATCH] 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>
+}
More information about the Mlir-commits
mailing list