Create a hybrid tf-tfl lowering pipline for TOSA

tensorflow/compiler/mlir/BUILD

@@ -111,6 +111,7 @@ cc_library(
         "//tensorflow/compiler/mlir/tensorflow:tf_saved_model_passes",
         "//tensorflow/compiler/mlir/tools/kernel_gen/ir:tf_framework_ops",
         "//tensorflow/compiler/mlir/tosa:tf_passes",
+        "//tensorflow/compiler/mlir/tosa:tf_tfl_passes",
         "//tensorflow/compiler/mlir/tosa:tfl_passes",
         "//tensorflow/compiler/mlir/xla:xla_legalize_tf",
         "//tensorflow/compiler/mlir/xla:xla_passes",

tensorflow/compiler/mlir/tf_mlir_opt_main.cc

@@ -29,6 +29,7 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.h"
 #include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
 #include "tensorflow/compiler/mlir/tosa/tf_passes.h"
+#include "tensorflow/compiler/mlir/tosa/tf_tfl_passes.h"
 #include "tensorflow/compiler/mlir/tosa/tfl_passes.h"
 #include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
 #include "tensorflow/compiler/mlir/xla/transforms/passes.h"
@@ -51,6 +52,7 @@ int main(int argc, char **argv) {
   mlir::tosa::registerLegalizeTosaPasses();
   mlir::tosa::registerTFtoTOSALegalizationPipeline();
   mlir::tosa::registerTFLtoTOSALegalizationPipeline();
+  mlir::tosa::registerTFTFLtoTOSALegalizationPipeline();
   mlir::tf_test::registerTensorFlowTestPasses();
   tensorflow::RegisterConvertMlirToXlaHloPipelineWithDefaults();
   tensorflow::RegisterGraphOptimizationPasses();

tensorflow/compiler/mlir/tosa/BUILD

@@ -188,3 +188,23 @@ cc_library(
         "@llvm-project//mlir:Transforms",
     ],
 )
+
+cc_library(
+    name = "tf_tfl_passes",
+    hdrs = [
+        "tf_tfl_passes.h",
+        "transforms/passes.h",
+    ],
+    srcs = [
+        "tf_tfl_passes.cc",
+        "transforms/legalize_tf_tfl.cc",
+    ],
+    compatible_with = get_compatible_with_cloud(),
+    visibility = [":friends"],
+    deps = [
+        ":passes_header",
+        ":tf_passes",
+        ":tfl_passes",
+        "//tensorflow/compiler/mlir/lite:lift_tflite_flex_ops",
+    ],
+)

tensorflow/compiler/mlir/tosa/tests/tf-tfl-to-tosa-pipeline.mlir

@@ -0,0 +1,23 @@
+// RUN: tf-opt --split-input-file --tf-tfl-to-tosa-pipeline  --verify-each %s | FileCheck %s
+
+// These tests focus on TensorFlow and TensorFlow Lite hybrid lowering. These
+// tests focus on tfl.custom operations that are Flex ops.
+
+// CHECK-LABEL: func @test_flex
+func @test_flex(%arg0: tensor<1x1x64xf32>) -> tensor<*xf32> {
+  %0 = "tfl.custom"(%arg0) {custom_code = "FlexIdentity", custom_option = opaque<"tfl", "0x084964656E74697479005212084964656E746974791A002A070A015412023001323B0A3966697273745F636F6E766F6C7574696F6E2F5446436F6E7631643178312F457870616E6444696D735F312F526561645661726961626C654F7000025E551414042801"> : tensor<102xi8>} : (tensor<1x1x64xf32>) -> tensor<*xf32>
+  // CHECK: %[[CAST:.+]] = tensor.cast %arg0
+  // CHECK: return %[[CAST]]
+  return %0 : tensor<*xf32>
+}
+
+// ----
+
+// CHECK-LABEL: func @test_flex_shape
+func @test_flex_shape(%arg0: tensor<?x2x64xf32>, %arg1: tensor<1x1x64xf32>) -> tensor<*xf32> {
+  // CHECK: %[[ADD:.+]] = "tosa.add"(%arg0, %arg1) : (tensor<?x2x64xf32>, tensor<1x1x64xf32>) -> tensor<?x2x64xf32>
+  // CHECK: %[[CAST:.+]] = tensor.cast %[[ADD]]
+  // CHECK: return %[[CAST]]
+  %0 = "tfl.custom"(%arg0, %arg1) {custom_code = "FlexAddV2", custom_option = opaque<"tfl", "0x054164645632002E120541646456321A001A002A070A01541202300132180A16726573696475616C5F626C6F636B5F2E5F302F616464000237311414042801"> : tensor<63xi8>} : (tensor<?x2x64xf32>, tensor<1x1x64xf32>) -> tensor<*xf32>
+  return %0: tensor<*xf32>
+}

tensorflow/compiler/mlir/tosa/tests/tf-to-tosa-pipeline.mlir

@@ -1,4 +1,5 @@
 // RUN: tf-opt --tf-to-tosa-pipeline  --verify-each %s | FileCheck %s
+// RUN: tf-opt --tf-tfl-to-tosa-pipeline  --verify-each %s | FileCheck %s
 
 // Operations for testing tf-to-tosa-pipeline
 // TODO: These tests are fairly minimal. Expand the checks to be more robust.
@@ -876,8 +877,7 @@ func @test_gather(%arg0: tensor<13x21x3xf32>) -> tensor<7x7x21x3xf32> {
 // CHECK-DAG: %[[VAR1:.*]] = "tosa.const"() {value = dense<1> : tensor<1x1xi32>}
 // CHECK-DAG: %[[VAR2:.*]] = "tosa.reshape"(%arg0) {new_shape = [1, 13, 63]}
 // CHECK-DAG: %[[VAR5:.*]] = "tosa.mul"(%[[VAR0]], %[[VAR1]]) {shift = 0 : i32}
-// CHECK-DAG: %[[VAR6:.*]] = "tosa.reduce_sum"(%[[VAR5]]) {axis = 1 : i64}
-// CHECK-DAG: %[[VAR7:.*]] = "tosa.reshape"(%[[VAR6]]) {new_shape = [1, 42]}
+// CHECK-DAG: %[[VAR7:.*]] = "tosa.reshape"(%[[VAR5]]) {new_shape = [1, 42]}
 // CHECK-DAG: %[[VAR8:.*]] = "tosa.gather"(%[[VAR2]], %[[VAR7]])
 // CHECK: %[[VAR9:.*]] = "tosa.reshape"(%[[VAR8]]) {new_shape = [6, 7, 21, 3]}
 func @test_gather_nd(%arg0: tensor<13x21x3xf32>) -> tensor<6x7x21x3xf32> {

tensorflow/compiler/mlir/tosa/tests/tfl-to-tosa-pipeline.mlir

@@ -1,4 +1,5 @@
 // RUN: tf-opt --split-input-file --tfl-to-tosa-pipeline --verify-each %s | FileCheck %s
+// RUN: tf-opt --split-input-file --tf-tfl-to-tosa-pipeline --verify-each %s | FileCheck %s
 
 // Operations for testing tfl-to-tosa-pipeline
 

tensorflow/compiler/mlir/tosa/tf_passes.cc

@@ -49,6 +49,7 @@ void createTFtoTOSALegalizationPipeline(
   //----------------------------------------------------------------------------
   // Post conversion cleanup.
   //----------------------------------------------------------------------------
+  pm.addPass(mlir::tosa::createTosaInferShapesPass());
   pm.addPass(mlir::tosa::createTosaMakeBroadcastablePass());
   // Inline the call/return basic blocks within TOSA control flow ops.
   pm.addPass(mlir::createInlinerPass());

tensorflow/compiler/mlir/tosa/tf_passes.h

@@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#ifndef TENSORFLOW_COMPILER_MLIR_TOSA_TOSA_PASSES_H
-#define TENSORFLOW_COMPILER_MLIR_TOSA_TOSA_PASSES_H
+#ifndef TENSORFLOW_COMPILER_MLIR_TOSA_TF_PASSES_H
+#define TENSORFLOW_COMPILER_MLIR_TOSA_TF_PASSES_H
 
 #include "mlir/Pass/PassManager.h"  // from @llvm-project
 #include "mlir/Pass/PassOptions.h"  // from @llvm-project
@@ -34,4 +34,4 @@ void registerTFtoTOSALegalizationPipeline();
 }  // namespace tosa
 }  // namespace mlir
 
-#endif  // TENSORFLOW_COMPILER_MLIR_TOSA_TOSA_PASSES_H
+#endif  // TENSORFLOW_COMPILER_MLIR_TOSA_TF_PASSES_H

tensorflow/compiler/mlir/tosa/tf_tfl_passes.cc

@@ -0,0 +1,73 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tosa/tf_tfl_passes.h"
+
+#include "mlir/Dialect/Affine/Passes.h"           // from @llvm-project
+#include "mlir/Dialect/Tosa/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"               // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/transforms/lift_tflite_flex_ops.h"
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+namespace mlir {
+namespace tosa {
+
+void createTFTFLtoTOSALegalizationPipeline(
+    OpPassManager& pm, const TOSATFTFLLegalizationPipelineOptions& opts) {
+  //----------------------------------------------------------------------------
+  // Prepare TFL module for conversion
+  //----------------------------------------------------------------------------
+  // Inline all functions into main and then delete the functions themselves.
+  pm.addPass(mlir::createInlinerPass());
+
+  // Add pass to decompose TFLite mixed quantization to non-quantized variants.
+  pm.addPass(TFL::CreateDecomposeHybridQuantizationPass());
+
+  // Now that there is only one function, run some MLIR passes on it.
+  pm.addPass(mlir::createCanonicalizerPass());
+  pm.addPass(mlir::createCSEPass());
+
+  pm.addPass(mlir::createLoopFusionPass());
+  pm.addPass(mlir::createAffineScalarReplacementPass());
+
+  //----------------------------------------------------------------------------
+  // Perform main conversion.
+  //----------------------------------------------------------------------------
+  pm.addPass(mlir::TFL::CreateLiftTfliteFlexOpsPass());
+  pm.addPass(mlir::tosa::createFuseBiasTFPass());
+  pm.addPass(mlir::tosa::createConvertTFLUint8Pass());
+  pm.addPass(mlir::tosa::createLegalizeTFTFLPass());
+
+  //----------------------------------------------------------------------------
+  // Post conversion cleanup.
+  //----------------------------------------------------------------------------
+  pm.addPass(mlir::tosa::createTosaInferShapesPass());
+  pm.addPass(mlir::tosa::createTosaMakeBroadcastablePass());
+  // Inline the call/return basic blocks within TOSA control flow ops.
+  pm.addPass(mlir::createInlinerPass());
+  // Clean up with DCE.
+  pm.addPass(mlir::createSymbolDCEPass());
+}
+
+void registerTFTFLtoTOSALegalizationPipeline() {
+  mlir::PassPipelineRegistration<TOSATFTFLLegalizationPipelineOptions>(
+      "tf-tfl-to-tosa-pipeline",
+      "TensorFlow / TensorFlow Lite to TOSA legalization pipeline",
+      createTFTFLtoTOSALegalizationPipeline);
+}
+
+}  // namespace tosa
+}  // namespace mlir

tensorflow/compiler/mlir/tosa/tf_tfl_passes.h

@@ -0,0 +1,37 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TOSA_TF_TFL_PASSES_H
+#define TENSORFLOW_COMPILER_MLIR_TOSA_TF_TFL_PASSES_H
+
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassOptions.h"  // from @llvm-project
+
+namespace mlir {
+namespace tosa {
+
+struct TOSATFTFLLegalizationPipelineOptions
+    : public PassPipelineOptions<TOSATFTFLLegalizationPipelineOptions> {};
+
+// Legalizes TF dialect(s) to Tosa.
+void createTFTFLtoTOSALegalizationPipeline(
+    OpPassManager& pm, const TOSATFTFLLegalizationPipelineOptions& opts);
+
+void registerTFTFLtoTOSALegalizationPipeline();
+
+}  // namespace tosa
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOSA_TF_TFL_PASSES_H

tensorflow/compiler/mlir/tosa/transforms/legalize_tf.cc

@@ -253,13 +253,13 @@ LogicalResult ConvertTFAddV2Op::matchAndRewrite(
     Operation* op, PatternRewriter& rewriter) const {
   auto tf_addv2_op = cast<TF::AddV2Op>(op);
 
-  RankedTensorType output_type =
-      tf_addv2_op.getResult().getType().dyn_cast<RankedTensorType>();
-  // Not a ranked tensor output
+  ShapedType output_type =
+      tf_addv2_op.getResult().getType().dyn_cast<ShapedType>();
+  // Not a shaped output
   if (!output_type) return failure();
 
-  rewriter.replaceOpWithNewOp<tosa::AddOp>(op, output_type, tf_addv2_op.x(),
-                                           tf_addv2_op.y());
+  CreateReplaceOpAndInfer<tosa::AddOp>(rewriter, op, output_type,
+                                       tf_addv2_op.x(), tf_addv2_op.y());
   return success();
 }
 
@@ -2272,10 +2272,17 @@ LogicalResult ConvertTFBatchMatMulV2Op::matchAndRewrite(
 }
 
 void LegalizeTF::runOnFunction() {
-  OwningRewritePatternList patterns(&getContext());
   auto* ctx = &getContext();
+  OwningRewritePatternList patterns(ctx);
   auto func = getFunction();
+  populateLegalizeTFPatterns(ctx, patterns);
+
+  ApplyPatternsWithShapeResolution(func, patterns);
+}
 
+}  // anonymous namespace
+
+void populateLegalizeTFPatterns(MLIRContext* ctx, RewritePatternSet& patterns) {
   // Add the generated patterns to the list.
   populateWithGenerated(patterns);
   patterns.insert<ConvertTFMatMulOp>(ctx);
@@ -2359,11 +2366,8 @@ void LegalizeTF::runOnFunction() {
   patterns.insert<ConvertTFRightShiftOp>(ctx);
   patterns.insert<ConvertTFOneHotOp>(ctx);
   patterns.insert<ConvertTFBatchMatMulV2Op>(ctx);
-  (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
 }
 
-}  // anonymous namespace
-
 // Creates an instance of the TensorFlow dialect LegalizeTF pass.
 std::unique_ptr<OperationPass<FuncOp>> createLegalizeTFPass() {
   return std::make_unique<LegalizeTF>();

tensorflow/compiler/mlir/tosa/transforms/legalize_tf_tfl.cc

@@ -0,0 +1,66 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Legalize TensorFlow and TensorFlow Lite to TOSA
+
+#include "mlir/Dialect/Quant/QuantOps.h"        // from @llvm-project
+#include "mlir/Dialect/Quant/QuantTypes.h"      // from @llvm-project
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"       // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"          // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"               // from @llvm-project
+#include "mlir/IR/PatternMatch.h"               // from @llvm-project
+#include "mlir/Support/LLVM.h"                  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_common.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_utils.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-legalize-tf-tfl"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir {
+namespace tosa {
+
+namespace {
+#define GEN_PASS_CLASSES
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+// Performs lowering to TOSA dialect
+class LegalizeTFTFL : public TosaLegalizeTFTFLPassBase<LegalizeTFTFL> {
+ public:
+  explicit LegalizeTFTFL() {}
+  void runOnFunction() override;
+};
+
+void LegalizeTFTFL::runOnFunction() {
+  auto* ctx = &getContext();
+  OwningRewritePatternList patterns(ctx);
+  auto func = getFunction();
+  populateLegalizeTFPatterns(ctx, patterns);
+  populateLegalizeTFLPatterns(ctx, patterns);
+
+  ApplyPatternsWithShapeResolution(func, patterns);
+}
+
+}  // anonymous namespace
+
+// Creates an instance of the TensorFlow Lite dialect LegalizeTFL pass.
+std::unique_ptr<OperationPass<FuncOp>> createLegalizeTFTFLPass() {
+  return std::make_unique<LegalizeTFTFL>();
+}
+
+}  // namespace tosa
+}  // namespace mlir