[TOSA] TorchToTosa option to emit partial conversion with Torch IR (#4106)

Add "require-full-tosa-conversion" option to TorchToTosa pipeline pass.
The default option is "true". When this option is set to "false", models
with non-legalized aten ops can still be partially converted with a
mixed of Torch and TOSA ops in the IR.

Example usage:

"builtin.module(torch-backend-to-tosa-backend-pipeline{require-full-tosa-conversion=false})"

Signed-off-by: Justin Ngo <justin.ngo@arm.com>

Author: justin-ngo-arm

include/torch-mlir/Conversion/Passes.td

@@ -122,6 +122,13 @@ def ConvertTorchToTosa : Pass<"convert-torch-to-tosa", "func::FuncOp"> {
     guards in case of shape mismatches.
   }];
   let constructor = "mlir::torch::createConvertTorchToTosaPass()";
+
+  let options = [
+    Option<"requireFullTosaConversion", "require-full-tosa-conversion",
+            "bool", /*default=*/"true",
+            "Require TorchToTosa full conversion by adding Torch Dialect to "
+            "TorchToTosa list of illegal dialects">,
+  ];
 }
 #endif
 

include/torch-mlir/Conversion/TorchToTosa/TorchToTosa.h

@@ -30,6 +30,8 @@ populateTorchToTosaConversionPatternsAndIllegalOps(TypeConverter &typeConverter,
                                                    RewritePatternSet &patterns);
 
 std::unique_ptr<OperationPass<func::FuncOp>> createConvertTorchToTosaPass();
+std::unique_ptr<OperationPass<func::FuncOp>>
+createConvertTorchToTosaPass(bool requireFullTosaConversion);
 } // namespace torch
 } // namespace mlir
 

include/torch-mlir/Dialect/TorchConversion/Transforms/Passes.h

@@ -28,9 +28,19 @@ void createTorchBackendToLinalgOnTensorsBackendPipeline(OpPassManager &pm);
 
 // Do not register the TOSA options if the TOSA target is disabled
 #ifdef TORCH_MLIR_ENABLE_TOSA
+struct TosaBackendPipelineOptions
+    : public PassPipelineOptions<TosaBackendPipelineOptions> {
+  Option<bool> requireFullTosaConversion{
+      *this, "require-full-tosa-conversion",
+      llvm::cl::desc("Require full TorchToTosa conversion by adding Torch "
+                     "Dialect to TorchToTosa list of illegal dialects"),
+      llvm::cl::init(true)};
+};
+
 /// Creates a pipeline that lowers from the torch backend contract to the
 /// TOSA backend contract.
-void createTorchBackendToTosaBackendPipeline(OpPassManager &pm);
+void createTorchBackendToTosaBackendPipeline(
+    OpPassManager &pm, const TosaBackendPipelineOptions &options);
 
 std::unique_ptr<OperationPass<ModuleOp>> createVerifyTosaBackendContractPass();
 #endif // TORCH_MLIR_ENABLE_TOSA

lib/Conversion/TorchToTosa/TorchToTosa.cpp

@@ -8950,6 +8950,11 @@ LogicalResult ConvertAtenOp<AtenUnfoldOp>::matchAndRewrite(
 namespace {
 class ConvertTorchToTosa : public ConvertTorchToTosaBase<ConvertTorchToTosa> {
 public:
+  ConvertTorchToTosa() = default;
+  ConvertTorchToTosa(bool requireFullTosaConversion) {
+    this->requireFullTosaConversion = requireFullTosaConversion;
+  }
+
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<tosa::TosaDialect>();
     registry.insert<tensor::TensorDialect>();
@@ -8962,7 +8967,12 @@ class ConvertTorchToTosa : public ConvertTorchToTosaBase<ConvertTorchToTosa> {
     ConversionTarget target(*context);
     target.addLegalDialect<tosa::TosaDialect, tensor::TensorDialect,
                            arith::ArithDialect>();
-    target.addIllegalDialect<Torch::TorchDialect>();
+
+    if (this->requireFullTosaConversion) {
+      target.addIllegalDialect<Torch::TorchDialect>();
+    } else {
+      target.addLegalDialect<Torch::TorchDialect>();
+    }
 
     TypeConverter typeConverter;
     typeConverter.addConversion([](Type type) { return type; });
@@ -9318,5 +9328,10 @@ std::set<StringRef> torch::populateTorchToTosaConversionPatternsAndIllegalOps(
 
 std::unique_ptr<OperationPass<func::FuncOp>>
 mlir::torch::createConvertTorchToTosaPass() {
-  return std::make_unique<ConvertTorchToTosa>();
+  return std::make_unique<ConvertTorchToTosa>(true);
+}
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+mlir::torch::createConvertTorchToTosaPass(bool requireFullTosaConversion) {
+  return std::make_unique<ConvertTorchToTosa>(requireFullTosaConversion);
 }

lib/Dialect/TorchConversion/Transforms/Passes.cpp

@@ -50,7 +50,7 @@ void mlir::torch::registerTorchConversionPasses() {
       "contract.",
       TorchConversion::createTorchBackendToLinalgOnTensorsBackendPipeline);
 #ifdef TORCH_MLIR_ENABLE_TOSA
-  mlir::PassPipelineRegistration<>(
+  mlir::PassPipelineRegistration<TorchConversion::TosaBackendPipelineOptions>(
       "torch-backend-to-tosa-backend-pipeline",
       "Pipeline lowering torch backend contract to TOSA backend "
       "contract.",
@@ -113,8 +113,10 @@ void TorchConversion::createTorchBackendToLinalgOnTensorsBackendPipeline(
 
 #ifdef TORCH_MLIR_ENABLE_TOSA
 void TorchConversion::createTorchBackendToTosaBackendPipeline(
-    OpPassManager &pm) {
-  pm.addNestedPass<func::FuncOp>(createConvertTorchToTosaPass());
+    OpPassManager &pm,
+    const TorchConversion::TosaBackendPipelineOptions &options) {
+  pm.addNestedPass<func::FuncOp>(
+      createConvertTorchToTosaPass(options.requireFullTosaConversion));
   // Perform rank broadcasting so TosaToLinalg pass works
   pm.addNestedPass<func::FuncOp>(createTosaMakeBroadcastablePass());