[mlir][math] expand-math pass assumes the static shaped type

[Lewuathe]

In the process of expand-math pass, the conversion of ceil op assumes the static shaped type as input as it needs create 0 and 1 constant values whose type is aligned with the op type.

Fixes #128275

[llvmbot]

@llvm/pr-subscribers-mlir-cf
@llvm/pr-subscribers-mlir-math

@llvm/pr-subscribers-mlir

Author: Kai Sasaki (Lewuathe)

Changes

In the process of expand-math pass, the conversion of ceil op assumes the static shaped type as input as it needs create 0 and 1 constant values whose type is aligned with the op type.

Fixes #128275

---

Full diff: https://github.com/llvm/llvm-project/pull/128299.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp (+5)
• (modified) mlir/test/Dialect/Math/expand-math.mlir (+22)

diff --git a/mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp b/mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp
index 23356d752146d..67e8dbba989b7 100644
--- a/mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/ExpandPatterns.cpp
@@ -222,6 +222,11 @@ 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 statis 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();
diff --git a/mlir/test/Dialect/Math/expand-math.mlir b/mlir/test/Dialect/Math/expand-math.mlir
index 1fdfb854325b4..4e249ec510afa 100644
--- a/mlir/test/Dialect/Math/expand-math.mlir
+++ b/mlir/test/Dialect/Math/expand-math.mlir
@@ -761,3 +761,25 @@ func.func @rsqrt_tns(%float: tensor<5x8xf32>) -> (tensor<5x8xf32>)  {
   %float_result = math.rsqrt %float : tensor<5x8xf32>
   return %float_result : tensor<5x8xf32>
 }
+
+// -----
+
+// CHECK-LABEL  func.func @non_static_shape_ceil_op
+// CHECK:     %[[IDX:.*]] = index.constant 0
+// CHECK:     %[[CST:.*]] = arith.constant dense<1.000000e+00> : tensor<2xf32>
+// CHECK:     %[[CAST:.*]] = tensor.cast %[[CST]] : tensor<2xf32> to tensor<?xf32>
+// CHECK:     %[[CEIL:.*]] = math.ceil %[[CAST]] : tensor<?xf32>
+// CHECK:     %[[DIM:.*]] = tensor.dim %[[CEIL]], %[[IDX]] : tensor<?xf32>
+// CHECK:     vector.print %[[DIM]] : index
+// CHECK:         return
+
+func.func @non_static_shape_ceil_op() {
+  %idx0 = index.constant 0
+  %cst_90 = arith.constant 1.000000e+00 : f32
+  %from_elements_92 = tensor.from_elements %cst_90, %cst_90 : tensor<2xf32>
+  %cast_93 = tensor.cast %from_elements_92 : tensor<2xf32> to tensor<?xf32>
+  %112 = math.ceil %cast_93 : tensor<?xf32>
+  %dim_233 = tensor.dim %112, %idx0 : tensor<?xf32>
+  vector.print %dim_233 : index
+  return
+}

[cferry-AMD]

Thanks! You may likewise also want to test unranked types:

func.func @unranked_ceil_op(%arg: tensor<*xf32>) -> tensor<*xf32>{
  %a = math.ceil %arg : tensor<*xf32>
  return %a: tensor<*xf32>
}

Also I see a constant created in rsqrt and we don't check for static types. I guess the same issue would apply there?

[Lewuathe]

Also I see a constant created in rsqrt and we don't check for static types. I guess the same issue would apply there?

Thanks for reporting that part. Let me take a look into the issue and will work on in another pull request if necessary.

[cferry-AMD]

LGTM. Thank you for this fix!