Refactor/cleanup hlo -> tfl pass

* Make the convert dot general function an explicit pattern class and register in cc
* Remove unused functions
* Remove redundant commentary

PiperOrigin-RevId: 649720051

Author: LukeBoyer

tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/dot_general.cc

@@ -15,6 +15,7 @@ limitations under the License.
 
 // This file implements logic for legalizing mhlo.dot_general to
 // tflite.batch_matmul.
+#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/dot_general.h"
 
 #include <cstddef>
 #include <cstdint>
@@ -39,6 +40,7 @@ limitations under the License.
 #include "mlir/IR/Value.h"  // from @llvm-project
 #include "mlir/IR/ValueRange.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
 #include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
 #include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
@@ -408,30 +410,15 @@ Value ConvertDot(PatternRewriter& rewriter, Value lhs, Value rhs,
   return reshaped.getResult();
 }
 
-// Converts mhlo.dot_general to tfl.BatchMatMul. Reshape and Transpose ops will
-// be inserted when necessary. See ConvertDotGeneralOp for additional notes.
-Value ConvertDotOp(PatternRewriter& rewriter, Operation* old_op) {
-  auto dot_op = cast<mhlo::DotOp>(old_op);
-  auto lhs_rank = mlir::cast<ShapedType>(dot_op.getLhs().getType()).getRank();
-  auto dot_dimension_numbers =
-      mhlo::DotDimensionNumbersAttr::get(rewriter.getContext(),
-                                         /*lhsBatchingDimensions=*/{},
-                                         /*rhsBatchingDimensions=*/{},
-                                         /*lhsContractingDimensions=*/
-                                         {lhs_rank == 1 ? 0 : 1},
-                                         /*rhsContractingDimensions=*/{0});
-  return ConvertDot(
-      rewriter, dot_op.getLhs(), dot_op.getRhs(), dot_dimension_numbers,
-      mlir::cast<ShapedType>(dot_op.getResult().getType()), dot_op.getLoc());
+LogicalResult LowerDotGeneralOp::matchAndRewrite(
+    mhlo::DotGeneralOp op, OpAdaptor adaptor,
+    ConversionPatternRewriter& rewriter) const {
+  auto val = ConvertDot(
+      rewriter, op.getLhs(), op.getRhs(), op.getDotDimensionNumbers(),
+      mlir::cast<ShapedType>(op.getResult().getType()), op.getLoc());
+  rewriter.replaceOp(op, val.getDefiningOp());
+  return mlir::success();
 }
 
-Value ConvertDotGeneralOp(PatternRewriter& rewriter, Operation* old_op) {
-  auto dot_general_op = cast<mhlo::DotGeneralOp>(old_op);
-  return ConvertDot(
-      rewriter, dot_general_op.getLhs(), dot_general_op.getRhs(),
-      dot_general_op.getDotDimensionNumbers(),
-      mlir::cast<ShapedType>(dot_general_op.getResult().getType()),
-      dot_general_op.getLoc());
-}
 }  // namespace odml
 }  // namespace mlir

tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/dot_general.h

@@ -18,12 +18,9 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_MLIR_LITE_STABLEHLO_TRANSFORMS_LEGALIZE_HLO_CONVERSIONS_DOT_GENERAL_H_
 #define TENSORFLOW_COMPILER_MLIR_LITE_STABLEHLO_TRANSFORMS_LEGALIZE_HLO_CONVERSIONS_DOT_GENERAL_H_
 
-#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
-#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
-#include "mlir/IR/Operation.h"  // from @llvm-project
-#include "mlir/IR/PatternMatch.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
 #include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
 
 namespace mlir {
 namespace odml {
@@ -41,7 +38,14 @@ namespace odml {
 // Concat ) are inserted for shape inference purposes.
 // 4) All the DotOp are converted to DotGeneral during the optimization pass
 // (ConvertDotOp).
-Value ConvertDotGeneralOp(PatternRewriter& rewriter, Operation* old_op);
+class LowerDotGeneralOp : public OpConversionPattern<mhlo::DotGeneralOp> {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      mhlo::DotGeneralOp op, OpAdaptor adaptor,
+      ConversionPatternRewriter& rewriter) const final;
+};
 }  // namespace odml
 }  // namespace mlir
 

tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo.cc

@@ -109,20 +109,22 @@ std::optional<bool> IsCbrtLegal(mhlo::CbrtOp op) {
 
 #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/generated_tflite_legalize_hlo.inc"
 void LegalizeHloToTfLitePass::runOnOperation() {
-  MLIRContext& context = getContext();
-  RewritePatternSet patterns(&getContext());
-  // Add new conversion patterns here.
-  PopulateLegalizeHloToTFLitePatterns(&patterns, &context);
+  MLIRContext* context = &getContext();
 
-  ConversionTarget target(context);
+  RewritePatternSet patterns(context);
+  patterns.add<odml::ConvertCustomCallOp, odml::LowerDotGeneralOp,
+               ConvertReduceOpToTFLiteArgmin, ConvertReduceOpToTFLiteArgmax>(
+      context);
+  populateWithGenerated(patterns);
+
+  ConversionTarget target(*context);
   target.addLegalDialect<TFL::TensorFlowLiteDialect, mhlo::MhloDialect>();
   target.addLegalOp<func::CallOp, func::ConstantOp, arith::ConstantOp>();
   target.addDynamicallyLegalOp<mhlo::CustomCallOp>(IsCustomCallLegal);
   target.addDynamicallyLegalOp<mhlo::ReduceOp>(IsReduceOpLegal);
-  // Converted MHLO ops should be marked illegal here.
-  // TODO: b/304003568 - Add TF_TransposeOp folding logic to tflite.
   target.addDynamicallyLegalOp<mhlo::CbrtOp>(IsCbrtLegal);
   target.addIllegalOp<mhlo::DotGeneralOp, mhlo::DotOp, mhlo::TransposeOp>();
+
   if (failed(applyPartialConversion(getOperation(), target,
                                     std::move(patterns)))) {
     getOperation().emitError("mhlo to TFLite legalization failed.");
@@ -131,14 +133,6 @@ void LegalizeHloToTfLitePass::runOnOperation() {
 }
 }  // namespace
 
-void PopulateLegalizeHloToTFLitePatterns(RewritePatternSet* patterns,
-                                         MLIRContext* context) {
-  patterns->add<odml::ConvertCustomCallOp>(context);
-  populateWithGenerated(*patterns);
-
-  patterns->add<ConvertReduceOpToTFLiteArgmin, ConvertReduceOpToTFLiteArgmax>(
-      context);
-}
 
 // Creates an instance of the pass.
 std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeHloToTfLitePass() {

tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo_patterns.td

@@ -13,8 +13,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-// This is the legalization pattern definition file for HLO to TFL.
-
 include "mlir/IR/OpBase.td"
 include "mlir/Dialect/Func/IR/FuncOps.td"
 include "mhlo/IR/hlo_ops.td"
@@ -23,30 +21,17 @@ include "tensorflow/compiler/mlir/lite/ir/tfl_ops.td"
 include "tensorflow/compiler/mlir/lite/ir/tfl_op_enums.td"
 include "mlir/Dialect/Arith/IR/ArithOps.td"
 
-
-
 def CreateTFLCastToInt32Op : NativeCodeCall<
   "CreateCastToInt32($0, $_loc, $_builder)">;
-// TODO: b/304003568 - Add TF_TransposeOp folding logic to tflite.
+
 def LegalizeTranspose : Pat<(MHLO_TransposeOp $arg, $perm),
                             (TFL_TransposeOp $arg,
                             (CreateTFLCastToInt32Op (TFL_ConstOp $perm)))>;
 
-
-def ConvertDotGeneralOp : NativeCodeCall<"ConvertDotGeneralOp($_builder, "
-                                               "$0.getDefiningOp())">;
-
-def LegalizeDotGeneral: Pat<(MHLO_DotGeneralOp:$old_value
-               $lhs,
-               $rhs,
-               $dot_dimension_numbers, $precision_config),
-          (ConvertDotGeneralOp $old_value)>;
-
-def LowerCbrt : Pat<(MHLO_CbrtOp $opr), 
+def LowerCbrt : Pat<(MHLO_CbrtOp $opr),
   (TFL_PowOp $opr,
     (TFL_DivOp
       (Arith_ConstantOp ConstantAttr<RankedF32ElementsAttr<[]>, "1.0f">),
       (Arith_ConstantOp ConstantAttr<RankedF32ElementsAttr<[]>, "3.0f">),
       TFL_AF_None)),
       [(F32Tensor $opr)]>;
-