[NFC][mlir][vector] Handle potential static cast assertion.

[shahidact]

In FoldArithToVectorOuterProduct pattern, static cast to vector type causes assertion when a scalar type was encountered. It seems the author meant to have a dyn_cast instead.

This NFC patch handles it by using dyn_cast.

[llvmbot]

@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Md Asghar Ahmad Shahid (shahidact)

Changes

In FoldArithToVectorOuterProduct pattern, static cast to vector type causes assertion when a scalar type was encountered. It seems the author meant to have a dyn_cast instead.

This NFC patch handles it by using dyn_cast.

---

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

1 Files Affected:

• (modified) mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp (+1-1)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 2269a40ec8ef1..023c4da7dffdf 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -2274,7 +2274,7 @@ struct FoldArithToVectorOuterProduct : public OpRewritePattern<MulOpType> {
 
   LogicalResult matchAndRewrite(MulOpType mulOp,
                                 PatternRewriter &rewriter) const override {
-    auto resType = llvm::cast<VectorType>(mulOp.getResult().getType());
+    auto resType = llvm::dyn_cast<VectorType>(mulOp.getResult().getType());
     if (!resType)
       return failure();
     if (resType.getRank() != 2)

[banach-space]

Thanks, do you have a repro that we could use as a test for this PR?

[shahidact]

Thanks, do you have a repro that we could use as a test for this PR?

Yes, I do have a test case, but I could not reproduce it using transform dialect while it is reproducible using a downstream pass.
Thats why I have not added any test.

[newling]

Looks like a good fix, but a test would be extra nice. Not possible to add one in the same file as the original PR?
#93664

Wondering if a version of
arith_to_outerproduct_trans_rhs_f32

which acts on tensor type instead of vector type would be sufficient?

[shahidact]

Looks like a good fix, but a test would be extra nice. Not possible to add one in the same file as the original PR?

#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
func.func @crash_test(
      %arg0: tensor<8x4xi8>, %arg1: tensor<4x8xi8>,
      %arg2: tensor<8x8xi32>) -> tensor<8x8xi32> {
    %0 = linalg.generic {
      indexing_maps = [#map, #map1, #map2],
      iterator_types = ["parallel", "parallel", "reduction"]}
      ins(%arg0, %arg1 : tensor<8x4xi8>, tensor<4x8xi8>)
      outs(%arg2 : tensor<8x8xi32>) {
    ^bb0(%in: i8, %in_0: i8, %out: i32):
      %0 = arith.extsi %in : i8 to i32
      %1 = arith.extsi %in_0 : i8 to i32
      %2 = arith.muli %0, %1 : i32
      %3 = arith.addi %out, %2 : i32
      linalg.yield %3 : i32
    } -> tensor<8x8xi32>
    return %0 : tensor<8x8xi32>
}

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
    %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
    %1 = transform.get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
    %2 = transform.structured.vectorize_children_and_apply_patterns %1  { fold_type_extensions_into_contract } : (!transform.any_op) -> !transform.any_op
    transform.yield
  }
}

I encountered the assert with above test case using the downstream vectorizer pass. I could not reproduce the assert with above transform schedule. Hence, I did not add the test case.

[newling]

Ok. This is clearly a bug, you're hitting it with a cast, it is resolved with a dyn_cast. So from my end it's not worth investing too much effort into getting a test. LGTM; thanks for the fix!

[joker-eph]

Ok. This is clearly a bug, you're hitting it with a cast, it is resolved with a dyn_cast. So from my end it's not worth investing too much effort into getting a test.

I am not sure why we wouldn't need a test actually?

Downstream often ncounter issues which are specific to their setup (they may be mis-using APIs, or have more test coverage), so the usual requirements is to have a test to provide coverage upstream.
So please follow-up with a test ASAP.

[newling]

I could not reproduce the assert with above transform schedule

I'm not very familiar with the transform dialect but my guess is that with your script it's not entering the block of the linalg.generic. I have added a test in #154434 that does hit the assert (before the fix in this PR).