[Mlir-commits] [mlir] add extra check on distribute types to avoid crashes (PR #102952)

[Bangtian Liu]

https://github.com/bangtianliu updated https://github.com/llvm/llvm-project/pull/102952

>From 3edb3003a35c066a9701eddf6e202a8440efa407 Mon Sep 17 00:00:00 2001
From: Bangtian Liu <liubangtian at gmail.com>
Date: Mon, 12 Aug 2024 11:55:00 -0700
Subject: [PATCH 1/2] add extra check on distribute types to avoid crashes

Signed-off-by: Bangtian Liu <liubangtian at gmail.com>
---
 mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
index 7285ad65fb549..29899f44eb2e2 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
@@ -1689,6 +1689,9 @@ struct WarpOpScfForOp : public OpRewritePattern<WarpExecuteOnLane0Op> {
           }
         });
 
+    if(llvm::any_of(distTypes, [](Type type){return !type;}))
+      return failure();
+
     SmallVector<size_t> newRetIndices;
     WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(
         rewriter, warpOp, escapingValues.getArrayRef(), distTypes,

>From 2a85257cf7b7e627f42e0c9cfe4bfc5451495e87 Mon Sep 17 00:00:00 2001
From: Bangtian Liu <liubangtian at gmail.com>
Date: Mon, 12 Aug 2024 12:12:15 -0700
Subject: [PATCH 2/2] format the code

Signed-off-by: Bangtian Liu <liubangtian at gmail.com>
---
 mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
index 29899f44eb2e2..6596f3fc0ed81 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
@@ -1689,7 +1689,7 @@ struct WarpOpScfForOp : public OpRewritePattern<WarpExecuteOnLane0Op> {
           }
         });
 
-    if(llvm::any_of(distTypes, [](Type type){return !type;}))
+    if (llvm::any_of(distTypes, [](Type type) { return !type; }))
       return failure();
 
     SmallVector<size_t> newRetIndices;