From c344695bca8f57c8e584fd7908144aead50d1f2a Mon Sep 17 00:00:00 2001
From: Stephan Herhut <herhut@google.com>
Date: Fri, 13 Nov 2020 00:57:54 -0800
Subject: [PATCH] Properly configure patterns in hlo to lhlo conversion with
 typeconverter.

Otherwise, materializations for partial conversions do not work.

PiperOrigin-RevId: 342215962
---
 lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc | 5 +----
 tests/end2end/broadcast.mlir                        | 7 ++++---
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc b/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc
index 59293e2..996bd10 100644
--- a/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc
+++ b/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc
@@ -555,11 +555,8 @@ struct HloLegalizeToLhlo
     ConversionTarget target(context);
     target.addLegalDialect<lmhlo::LmhloDialect>();
     target.addLegalDialect<StandardOpsDialect>();
-    target.addLegalOp<ModuleOp>();
     target.addIllegalOp<mlir::TensorLoadOp>();
     target.addIllegalOp<mlir::TensorStoreOp>();
-    target.addLegalOp<ModuleTerminatorOp>();
-    target.addLegalOp<TensorFromElementsOp>();
     target.addIllegalDialect<mhlo::MhloDialect>();
 
     BufferizeTypeConverter converter;
@@ -645,7 +642,7 @@ void populateHLOToLHLOConversionPattern(MLIRContext* context,
       HloToLhloReturnOpConverter,
       HloToLhloTensorLoadOpConverter,
       HloToLhloTensorStoreOpConverter
-  >(context);
+  >(*converter, context);
   // clang-format on
 }
 
diff --git a/tests/end2end/broadcast.mlir b/tests/end2end/broadcast.mlir
index 9d9a2df..544efb8 100644
--- a/tests/end2end/broadcast.mlir
+++ b/tests/end2end/broadcast.mlir
@@ -1,7 +1,8 @@
 // RUN: mlir-hlo-opt %s -chlo-legalize-to-hlo -hlo-legalize-to-lhlo \
-// RUN: -buffer-hoisting -buffer-deallocation -copy-removal -canonicalize -cse \
-// RUN: -lhlo-legalize-to-linalg -lhlo-fuse-linalg -convert-linalg-to-loops \
-// RUN: -canonicalize -cse -convert-linalg-to-llvm -convert-std-to-llvm \
+// RUN: -std-bufferize --canonicalize -buffer-hoisting -buffer-deallocation \
+// RUN: -copy-removal -canonicalize -cse -lhlo-legalize-to-linalg \
+// RUN: -lhlo-fuse-linalg -convert-linalg-to-loops -canonicalize -cse \
+// RUN: -convert-linalg-to-llvm -convert-std-to-llvm \
 // RUN: | mlir-cpu-runner -e main -entry-point-result=void \
 // RUN: -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s --dump-input=always