[mlir][arith] Only fold splats for static shape result types

[KB9]

This prevents an assertion when constructing the DenseElementsAttr result, where the passed-in type is expected to have a static shape.

Fixes #92057

[github-actions]

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be
notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write
permissions for the repository. In which case you can instead tag reviewers by
name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review
by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate
is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-arith

Author: Kavan Bickerstaff (KB9)

Changes

This prevents an assertion when constructing the DenseElementsAttr result, where the passed-in type is expected to have a static shape.

Fixes #92057

---

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

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/CommonFolders.h (+4-1)
• (modified) mlir/test/Dialect/Arith/canonicalize.mlir (+8)

diff --git a/mlir/include/mlir/Dialect/CommonFolders.h b/mlir/include/mlir/Dialect/CommonFolders.h
index 7dabc781cd595..6f497a259262a 100644
--- a/mlir/include/mlir/Dialect/CommonFolders.h
+++ b/mlir/include/mlir/Dialect/CommonFolders.h
@@ -298,7 +298,10 @@ Attribute constFoldCastOp(ArrayRef<Attribute> operands, Type resType,
         calculate(op.getSplatValue<ElementValueT>(), castStatus);
     if (!castStatus)
       return {};
-    return DenseElementsAttr::get(cast<ShapedType>(resType), elementResult);
+    auto shapedResType = cast<ShapedType>(resType);
+    if (!shapedResType.hasStaticShape())
+      return {};
+    return DenseElementsAttr::get(shapedResType, elementResult);
   }
   if (auto op = dyn_cast<ElementsAttr>(operands[0])) {
     // Operand is ElementsAttr-derived; perform an element-wise fold by
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index e4f95bb0545a2..1a387c20c4b29 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -2950,6 +2950,14 @@ func.func @unsignedExtendConstantResource() -> tensor<i16> {
   return %ext : tensor<i16>
 }
 
+// Just checks that this doesn't crash.
+// CHECK-LABEL: @signedExtendSplatAsDynamicShape
+func.func @signedExtendSplatAsDynamicShape() -> tensor<?xi64> {
+  %splat = arith.constant dense<5> : tensor<2xi16>
+  %extsplat = arith.extsi %splat : tensor<2xi16> to tensor<?xi64>
+  return %extsplat : tensor<?xi64>
+}
+
 // CHECK-LABEL: @extsi_i0
 //       CHECK:   %[[ZERO:.*]] = arith.constant 0 : i16
 //       CHECK:   return %[[ZERO]] : i16

[MaheshRavishankar]

Thanks!

[kuhar]

I'm surprised this is even allowed at the level of artith. Shouldn't we change the verifier to reject this extend + 'shape cast' pattern instead?

[KB9]

Makes sense. I'll rewrite it so that it's rejected during verification.

[MaheshRavishankar]

That might be a different can of worms. This is an obvious things to fix for now. Maybe we should land this first.

[kuhar]

Currently the consensus seems to be to disallow these casts (i.e., element type and static -> dynamic) in arith: https://discourse.llvm.org/t/rfc-remove-arith-math-ops-on-tensors/74357/68 .

But I'm fine with this as a stop-gap until that happens.

[KB9]

Thanks for the reviews! Someone else will need to merge this for me, as I don't appear to have write access to do so.

[kuhar]

I can merge it for you, but it seems like you don't have a public email on github. This is strongly preferred so that buildbots / people can reach you when you PR is discovered to cause issues.

[KB9]

Thanks for the advice. I've made my email public.

[github-actions]

@KB9 Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested
by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as
the builds can include changes from many authors. It is not uncommon for your
change to be included in a build that fails due to someone else's changes, or
infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself.
This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!