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

[Bangtian Liu]

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

This PR addresses the issue detailed in https://github.com/iree-org/iree/issues/17948. 

The problem occurs when distributed types are set to NULL, leading to compilation crashes.









>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] 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 7285ad65fb549e..29899f44eb2e22 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,