From 61244b136cfdbb0ba6459d06df4f27d3bfbd0887 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Wed, 16 Dec 2020 04:57:35 -0800
Subject: [PATCH] Try to avoid a segfault if we don't support a lowering.

It can happen that a lowering for a certain type is not implemented yet.
We should not segfault in such a case, but instead return a failure().

PiperOrigin-RevId: 347801106
---
 lib/Dialect/mhlo/transforms/legalize_to_linalg.cc | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc b/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
index cba8a52..4d44b9c 100644
--- a/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
+++ b/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
@@ -131,6 +131,7 @@ class PointwiseToLinalgConverter : public OpConversionPattern<OpTy> {
     SmallVector<AffineMap, 2> indexing_maps(args.size() + (isLHLO ? 0 : 1),
                                             common_indexing_map);
 
+    bool failed = false;
     auto linalg_op = rewriter.create<linalg::GenericOp>(
         loc, op_result_types, inputs, output_buffers,
         /*initTensors=*/ValueRange{}, indexing_maps,
@@ -141,8 +142,13 @@ class PointwiseToLinalgConverter : public OpConversionPattern<OpTy> {
           Value op_result = lmhlo::HloOpToStdScalarOp::map<OpTy>(
               op, body_result_types,
               llvm::to_vector<2>(args.take_front(inputs.size())), &rewriter);
-          nested_builder.create<linalg::YieldOp>(loc, op_result);
+          if (op_result == nullptr) {
+            failed = true;
+          } else {
+            nested_builder.create<linalg::YieldOp>(loc, op_result);
+          }
         });
+    if (failed) return failure();
     rewriter.replaceOp(op, linalg_op.getOperation()->getResults());
     return success();
   }