[mlir] Fix vectorize transform crashing on none-op operand

diff --git a/mlir/lib/Transforms/Vectorize.cpp b/mlir/lib/Transforms/Vectorize.cpp
index 9b64311b4192fd5c3349b00348a65ea8a6eb97bd..617a7537bcca1410870e0a14e47923db485b41cc 100644 (file)
--- a/mlir/lib/Transforms/Vectorize.cpp
+++ b/mlir/lib/Transforms/Vectorize.cpp
@@ -1003,7 +1003,7 @@ static Value vectorizeOperand(Value operand, Operation *op,
     return nullptr;
   }
   // 3. vectorize constant.
-  if (auto constant = dyn_cast<ConstantOp>(operand.getDefiningOp())) {
+  if (auto constant = dyn_cast_or_null<ConstantOp>(operand.getDefiningOp())) {
     return vectorizeConstant(
         op, constant,
         VectorType::get(state->strategy->vectorSizes, operand.getType()));
@@ -1037,6 +1037,8 @@ static Operation *vectorizeOneOperation(Operation *opInst,
     auto memRef = store.getMemRef();
     auto value = store.getValueToStore();
     auto vectorValue = vectorizeOperand(value, opInst, state);
+    if (!vectorValue)
+      return nullptr;
 
     ValueRange mapOperands = store.getMapOperands();
     SmallVector<Value, 8> indices;