[MLIR] More graceful failure in MaterializeVectors

Even though it is unexpected except in pathological cases, a nullptr clone may
be returned. This CL handles the nullptr return gracefuly.

PiperOrigin-RevId: 227764615

1 files changed, 6 insertions, 3 deletions

diff --git a/mlir/lib/Transforms/MaterializeVectors.cpp b/mlir/lib/Transforms/MaterializeVectors.cpp
index bbe2f85319b..e687889cf93 100644
--- a/mlir/lib/Transforms/MaterializeVectors.cpp
+++ b/mlir/lib/Transforms/MaterializeVectors.cpp
@@ -554,12 +554,15 @@ static bool instantiateMaterialization(Instruction *inst,
   FuncBuilder b(inst);
   auto *opInst = cast<OperationInst>(inst);
   if (auto write = opInst->dyn_cast<VectorTransferWriteOp>()) {
-    instantiate(&b, write, state->hwVectorType, state->hwVectorInstance,
-                state->substitutionsMap);
-    return false;
+    auto *clone = instantiate(&b, write, state->hwVectorType,
+                              state->hwVectorInstance, state->substitutionsMap);
+    return clone == nullptr;
   } else if (auto read = opInst->dyn_cast<VectorTransferReadOp>()) {
     auto *clone = instantiate(&b, read, state->hwVectorType,
                               state->hwVectorInstance, state->substitutionsMap);
+    if (!clone) {
+      return true;
+    }
     state->substitutionsMap->insert(
         std::make_pair(read->getResult(), clone->getResult(0)));
     return false;