commit 15c26be8f79834ae5a6123939b4885e4d6a112af
author Eugene Zhulenev <ezhulenev@google.com> Mon May 23 21:21:22 2022 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Mon May 23 21:24:41 2022 -0700
tree 17d8f90dabebb0ef8a7b7f2d96ed43800bd31e88
parent eee70a365f5deb9c78e63945756f856fb96847cb

[xla:jitrt] Work around the canonicalization bug for mhlo copy operation

mhlo.copy operation canonicalizer can break downstream gpu code generation, for now just make sure that the first pass only removes memref.get_global ops

PiperOrigin-RevId: 450593603

tensorflow/compiler/mlir/tfrt/transforms/lmhlo_to_gpu/lmhlo_to_jitrt.cc

diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lmhlo_to_gpu/lmhlo_to_jitrt.cc b/tensorflow/compiler/mlir/tfrt/transforms/lmhlo_to_gpu/lmhlo_to_jitrt.cc
index 3e0771f..7095464 100644
--- a/tensorflow/compiler/mlir/tfrt/transforms/lmhlo_to_gpu/lmhlo_to_jitrt.cc
+++ b/tensorflow/compiler/mlir/tfrt/transforms/lmhlo_to_gpu/lmhlo_to_jitrt.cc
@@ -469,6 +469,11 @@
   target.addIllegalOp<GetGlobalOp>();
   target.addLegalOp<ConstantOp, memref::ViewOp>();
 
+  // TODO(ezhulenev): By adding MHLO and LMHLO to a set of legal dialects, we
+  // suppress any rewrites for these dialects (there are canonicalization
+  // patterns that interact badly with downstream Gpu binary code generation).
+  target.addLegalDialect<mhlo::MhloDialect, lmhlo::LmhloDialect>();
+
   if (failed(applyPartialConversion(module, target, std::move(patterns))))
     signalPassFailure();
 }