Add a pass pipeline to convert Torch to a mix of Tosa, Linalg, Tensor and other standard MLIR dialects.

sahas3 · sahas3 · commit fa9d38518be3 · 2025-07-10T13:06:36.000-04:00
diff --git a/include/torch-mlir/Conversion/Passes.td b/include/torch-mlir/Conversion/Passes.td
@@ -132,6 +132,14 @@ def ConvertTorchToTosa : Pass<"convert-torch-to-tosa", "func::FuncOp"> {
 }
 #endif
 
+def ConvertTorchToTosaLinalg : Pass<"convert-torch-to-tosa-linalg", "func::FuncOp"> {
+  let summary = "Convert Torch ops to a mix of TOSA ops and LINALG_ON_TENSORS ops";
+  let description = [{
+    This pass tries to lower torch ops to tosa ops if possible. Otherwise lowers to a mix of linalg, tensor, scf, and other standard mlir dialects.
+  }];
+  let constructor = "mlir::torch::createConvertTorchToTosaLinalgPass()";
+}
+
 def ConvertTorchToTMTensor : Pass<"convert-torch-to-tmtensor", "func::FuncOp"> {
   let summary = "Convert recognized Torch ops to TMTensor/Linalg ops";
   let description = [{
diff --git a/include/torch-mlir/Conversion/TorchToLinalg/TorchToLinalg.h b/include/torch-mlir/Conversion/TorchToLinalg/TorchToLinalg.h
@@ -13,12 +13,18 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
 #include <memory>
 
 namespace mlir {
 namespace torch {
+
+void populateTorchToLinalgOnTensorsPatternsAndLegality(
+    TypeConverter &typeConverter, RewritePatternSet &patterns,
+    ConversionTarget &target);
+
 std::unique_ptr<OperationPass<func::FuncOp>> createConvertTorchToLinalgPass();
-}
+} // namespace torch
 } // namespace mlir
 
 #endif // TORCHMLIR_CONVERSION_ATENTOLINALG_ATENTOLINALG_H
diff --git a/include/torch-mlir/Conversion/TorchToTosaLinalg/TorchToTosaLinalg.h b/include/torch-mlir/Conversion/TorchToTosaLinalg/TorchToTosaLinalg.h
@@ -0,0 +1,24 @@
+//===------------------------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// Also available under a BSD-style license. See LICENSE.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TORCHMLIR_CONVERSION_ATENTOTOSALINALG_ATENTOTOSALINALG_H
+#define TORCHMLIR_CONVERSION_ATENTOTOSALINALG_ATENTOTOSALINALG_H
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Pass/Pass.h"
+#include <memory>
+
+namespace mlir {
+namespace torch {
+std::unique_ptr<OperationPass<func::FuncOp>>
+createConvertTorchToTosaLinalgPass();
+}
+} // namespace mlir
+
+#endif // TORCHMLIR_CONVERSION_ATENTOTOSALINALG_ATENTOTOSALINALG_H
diff --git a/include/torch-mlir/Dialect/TorchConversion/Transforms/Passes.h b/include/torch-mlir/Dialect/TorchConversion/Transforms/Passes.h
@@ -45,6 +45,10 @@ void createTorchBackendToTosaBackendPipeline(
 std::unique_ptr<OperationPass<ModuleOp>> createVerifyTosaBackendContractPass();
 #endif // TORCH_MLIR_ENABLE_TOSA
 
+/// Creates a pipeline that lowers from the torch backend contract to the
+/// TOSA + linalg backend contract.
+void createTorchBackendToTosaLinalgBackendPipeline(OpPassManager &pm);
+
 // Do not register the stablehlo options if the stablehlo target is disabled
 #ifdef TORCH_MLIR_ENABLE_STABLEHLO
 struct StablehloBackendPipelineOptions
@@ -92,6 +96,9 @@ createConvertCustomQuantOpPass();
 std::unique_ptr<OperationPass<ModuleOp>>
 createVerifyLinalgOnTensorsBackendContractPass();
 
+std::unique_ptr<OperationPass<ModuleOp>>
+createVerifyTosaLinalgBackendContractPass();
+
 } // namespace TorchConversion
 
 /// Registers all Torch transformation passes.
diff --git a/include/torch-mlir/Dialect/TorchConversion/Transforms/Passes.td b/include/torch-mlir/Dialect/TorchConversion/Transforms/Passes.td
@@ -68,6 +68,11 @@ def VerifyTosaBackendContract : Pass<"torch-verify-tosa-backend-contract", "Modu
 }
 #endif
 
+def VerifyTosaLinalgBackendContract : Pass<"torch-verify-tosa-linalg-backend-contract", "ModuleOp"> {
+  let summary = "Verifies conformity to the tosa + linalg-on-tensors backend contract";
+  let constructor = "mlir::torch::TorchConversion::createVerifyTosaLinalgBackendContractPass()";
+}
+
 #ifdef TORCH_MLIR_ENABLE_STABLEHLO
 def VerifyStablehloBackendContract : Pass<"torch-verify-stablehlo-backend-contract", "ModuleOp"> {
   let summary = "Verifies conformity to the stablehlo backend contract";
diff --git a/lib/Conversion/CMakeLists.txt b/lib/Conversion/CMakeLists.txt
@@ -5,6 +5,7 @@ add_subdirectory(TorchToSCF)
 add_subdirectory(TorchToTensor)
 if(TORCH_MLIR_ENABLE_TOSA)
   add_subdirectory(TorchToTosa)
+  add_subdirectory(TorchToTosaLinalg)
 endif()
 if(TORCH_MLIR_ENABLE_STABLEHLO)
   add_subdirectory(TorchToStablehlo)
@@ -25,7 +26,7 @@ if(TORCH_MLIR_ENABLE_STABLEHLO)
   list(APPEND linked_libs TorchMLIRTorchToStablehlo)
 endif()
 if(TORCH_MLIR_ENABLE_TOSA)
-  list(APPEND linked_libs TorchMLIRTorchToTosa)
+  list(APPEND linked_libs TorchMLIRTorchToTosa TorchMLIRTorchToTosaLinalg)
 endif()
 
 add_mlir_library(TorchMLIRConversionPasses
diff --git a/lib/Conversion/Passes.cpp b/lib/Conversion/Passes.cpp
@@ -22,6 +22,7 @@
 
 #ifdef TORCH_MLIR_ENABLE_TOSA
 #include "torch-mlir/Conversion/TorchToTosa/TorchToTosa.h"
+#include "torch-mlir/Conversion/TorchToTosaLinalg/TorchToTosaLinalg.h"
 #endif // TORCH_MLIR_ENABLE_TOSA
 
 //===----------------------------------------------------------------------===//
diff --git a/lib/Conversion/TorchToLinalg/TorchToLinalg.cpp b/lib/Conversion/TorchToLinalg/TorchToLinalg.cpp
@@ -63,24 +63,8 @@ class ConvertTorchToLinalg
 
     RewritePatternSet patterns(context);
 
-    torch_to_linalg::populateTensorScalarInteropPatternsAndLegality(
-        typeConverter, patterns, target);
-    torch_to_linalg::populateLinearPatternsAndLegality(typeConverter, patterns,
-                                                       target);
-    torch_to_linalg::populatePoolingPatternsAndLegality(typeConverter, patterns,
-                                                        target);
-    torch_to_linalg::populateRandomPatternsAndLegality(typeConverter, patterns,
-                                                       target);
-    torch_to_linalg::populateUncategorizedPatternsAndLegality(typeConverter,
-                                                              patterns, target);
-    torch_to_linalg::populateReductionPatternsAndLegality(typeConverter,
-                                                          patterns, target);
-    torch_to_linalg::populateDataMovementPatternsAndLegality(typeConverter,
+    torch::populateTorchToLinalgOnTensorsPatternsAndLegality(typeConverter,
                                                              patterns, target);
-    torch_to_linalg::populateIndirectDataMovementPatternsAndLegality(
-        typeConverter, patterns, target);
-    torch_to_linalg::populateTensorConstructorsPatternsAndLegality(
-        typeConverter, patterns, target);
 
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
@@ -89,6 +73,29 @@ class ConvertTorchToLinalg
 };
 } // namespace
 
+void mlir::torch::populateTorchToLinalgOnTensorsPatternsAndLegality(
+    TypeConverter &typeConverter, RewritePatternSet &patterns,
+    ConversionTarget &target) {
+  torch_to_linalg::populateTensorScalarInteropPatternsAndLegality(
+      typeConverter, patterns, target);
+  torch_to_linalg::populateLinearPatternsAndLegality(typeConverter, patterns,
+                                                     target);
+  torch_to_linalg::populatePoolingPatternsAndLegality(typeConverter, patterns,
+                                                      target);
+  torch_to_linalg::populateRandomPatternsAndLegality(typeConverter, patterns,
+                                                     target);
+  torch_to_linalg::populateUncategorizedPatternsAndLegality(typeConverter,
+                                                            patterns, target);
+  torch_to_linalg::populateReductionPatternsAndLegality(typeConverter, patterns,
+                                                        target);
+  torch_to_linalg::populateDataMovementPatternsAndLegality(typeConverter,
+                                                           patterns, target);
+  torch_to_linalg::populateIndirectDataMovementPatternsAndLegality(
+      typeConverter, patterns, target);
+  torch_to_linalg::populateTensorConstructorsPatternsAndLegality(
+      typeConverter, patterns, target);
+}
+
 std::unique_ptr<OperationPass<func::FuncOp>>
 mlir::torch::createConvertTorchToLinalgPass() {
   return std::make_unique<ConvertTorchToLinalg>();
diff --git a/lib/Conversion/TorchToTosaLinalg/CMakeLists.txt b/lib/Conversion/TorchToTosaLinalg/CMakeLists.txt
@@ -0,0 +1,19 @@
+add_mlir_conversion_library(TorchMLIRTorchToTosaLinalg
+  TorchToTosaLinalg.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${PROJECT_SOURCE_DIR}/include/torch-mlir/Conversion/TorchToTosaLinalg
+
+  DEPENDS
+  TorchMLIRConversionPassIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRPass
+  MLIRTosaDialect
+  MLIRLinalgDialect
+  TorchMLIRConversionUtils
+  TorchMLIRTorchDialect
+)
+
+torch_mlir_target_includes(TorchMLIRTorchToTosaLinalg)
diff --git a/lib/Conversion/TorchToTosaLinalg/TorchToTosaLinalg.cpp b/lib/Conversion/TorchToTosaLinalg/TorchToTosaLinalg.cpp
@@ -0,0 +1,86 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// Also available under a BSD-style license. See LICENSE.
+//
+//===----------------------------------------------------------------------===//
+
+#include "torch-mlir/Conversion/TorchToTosaLinalg/TorchToTosaLinalg.h"
+#include "torch-mlir/Conversion/TorchToLinalg/TorchToLinalg.h"
+#include "torch-mlir/Conversion/TorchToTosa/TorchToTosa.h"
+
+#include "../PassDetail.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "torch-mlir/Dialect/Torch/IR/TorchDialect.h"
+#include "torch-mlir/Dialect/Torch/IR/TorchOps.h"
+#include "torch-mlir/Dialect/TorchConversion/IR/TorchConversionOps.h"
+#include "torch-mlir/Dialect/TorchConversion/Transforms/BackendTypeConversion.h"
+
+using namespace mlir;
+using namespace mlir::torch;
+using namespace mlir::torch::Torch;
+
+// -----------------------------------------------------------------------------
+// The pass
+// -----------------------------------------------------------------------------
+
+namespace {
+class ConvertTorchToTosaLinalg
+    : public ConvertTorchToTosaLinalgBase<ConvertTorchToTosaLinalg> {
+public:
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<linalg::LinalgDialect>();
+    registry.insert<math::MathDialect>();
+    registry.insert<func::FuncDialect>();
+    registry.insert<tensor::TensorDialect>();
+    registry.insert<tosa::TosaDialect>();
+    registry.insert<arith::ArithDialect>();
+    registry.insert<cf::ControlFlowDialect>();
+    registry.insert<scf::SCFDialect>();
+    registry.insert<complex::ComplexDialect>();
+    TorchConversion::getBackendTypeConversionDependentDialects(registry);
+  }
+
+  void runOnOperation() override {
+    MLIRContext *context = &getContext();
+    ConversionTarget target(*context);
+    target.addLegalDialect<linalg::LinalgDialect, func::FuncDialect,
+                           cf::ControlFlowDialect, math::MathDialect,
+                           scf::SCFDialect, sparse_tensor::SparseTensorDialect,
+                           tosa::TosaDialect, tensor::TensorDialect,
+                           arith::ArithDialect, complex::ComplexDialect>();
+    target.addLegalOp<TorchConversion::GetNextSeedOp>();
+    torch::populateTorchToTosaConversionLegalOps(target);
+
+    TypeConverter typeConverter;
+    typeConverter.addConversion([](Type type) { return type; });
+    TorchConversion::setupBackendTypeConversion(target, typeConverter);
+
+    RewritePatternSet patterns(context);
+
+    torch::populateTorchToTosaConversionPatterns(typeConverter, patterns);
+    torch::populateTorchToLinalgOnTensorsPatternsAndLegality(typeConverter,
+                                                             patterns, target);
+
+    if (failed(applyPartialConversion(getOperation(), target,
+                                      std::move(patterns))))
+      return signalPassFailure();
+  }
+};
+
+} // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+mlir::torch::createConvertTorchToTosaLinalgPass() {
+  return std::make_unique<ConvertTorchToTosaLinalg>();
+}
diff --git a/lib/Dialect/TorchConversion/Transforms/CMakeLists.txt b/lib/Dialect/TorchConversion/Transforms/CMakeLists.txt
@@ -27,6 +27,7 @@ add_mlir_library(TorchMLIRTorchConversionPasses
   UnpackQuantTensor.cpp
   VerifyLinalgOnTensorsBackendContract.cpp
   VerifyTosaBackendContract.cpp
+  VerifyTosaLinalgBackendContract.cpp
   VerifyStablehloBackendContract.cpp
 
   ADDITIONAL_HEADER_DIRS
diff --git a/lib/Dialect/TorchConversion/Transforms/Passes.cpp b/lib/Dialect/TorchConversion/Transforms/Passes.cpp
@@ -18,6 +18,8 @@
 #include "torch-mlir/Conversion/TorchToSCF/TorchToSCF.h"
 #include "torch-mlir/Conversion/TorchToTMTensor/TorchToTMTensor.h"
 #include "torch-mlir/Conversion/TorchToTensor/TorchToTensor.h"
+#include "torch-mlir/Conversion/TorchToTosa/TorchToTosa.h"
+#include "torch-mlir/Conversion/TorchToTosaLinalg/TorchToTosaLinalg.h"
 #include "torch-mlir/Dialect/Torch/Transforms/Passes.h"
 
 #ifdef TORCH_MLIR_ENABLE_STABLEHLO
@@ -55,7 +57,14 @@ void mlir::torch::registerTorchConversionPasses() {
       "Pipeline lowering torch backend contract to TOSA backend "
       "contract.",
       TorchConversion::createTorchBackendToTosaBackendPipeline);
+
+  mlir::PassPipelineRegistration<>(
+      "torch-backend-to-tosa-linalg-backend-pipeline",
+      "Pipeline lowering torch backend contract to TOSA + Linalg backend "
+      "contract.",
+      TorchConversion::createTorchBackendToTosaLinalgBackendPipeline);
 #endif
+
 #ifdef TORCH_MLIR_ENABLE_STABLEHLO
   mlir::PassPipelineRegistration<
       TorchConversion::StablehloBackendPipelineOptions>(
@@ -143,6 +152,58 @@ void TorchConversion::createTorchBackendToTosaBackendPipeline(
   // correct form.
   pm.addPass(TorchConversion::createVerifyTosaBackendContractPass());
 }
+
+void TorchConversion::createTorchBackendToTosaLinalgBackendPipeline(
+    OpPassManager &pm) {
+
+  // We want to fuse quantized operations together before lowering to
+  // linalg.
+  pm.addNestedPass<func::FuncOp>(Torch::createFuseQuantizedOpsPass());
+  pm.addNestedPass<func::FuncOp>(Torch::createScalarizeShapesPass());
+
+  // Lower to tosa+linalg + guards which is the input to codegen backends.
+  // We do this first as it tends to involve pattern-matching against constants,
+  // (e.g. dimensions which must be constant in a ranked programming model)
+  // and those constants get somewhat obscured by TorchToArith.
+  pm.addNestedPass<func::FuncOp>(createConvertTorchToTMTensorPass());
+  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+
+  // Tosa+Linalg pipeline pass
+  pm.addNestedPass<func::FuncOp>(createConvertTorchToTosaLinalgPass());
+  // Perform rank broadcasting so TosaToLinalg pass works
+  pm.addNestedPass<func::FuncOp>(createTosaMakeBroadcastablePass());
+
+  // Clean up any non-canonical code introduced above..
+  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+  // The resolution of `dim` ops tends to create identical ops. CSE them.
+  pm.addNestedPass<func::FuncOp>(createCSEPass());
+  pm.addNestedPass<func::FuncOp>(createConvertTorchToSCFPass());
+  pm.addNestedPass<func::FuncOp>(createConvertTorchToArithPass());
+  pm.addNestedPass<func::FuncOp>(createConvertTorchToTensorPass());
+  pm.addPass(createConvertTorchConversionToMLProgramPass());
+  pm.addNestedPass<func::FuncOp>(memref::createExpandOpsPass());
+
+  // Clean up any non-canonical code introduced above..
+  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+  // Resolve `dim` ops on tensors (which currently live in the `memref`
+  // dialect for some reason -- we don't have memrefs at this level).
+  pm.addNestedPass<func::FuncOp>(
+      memref::createResolveShapedTypeResultDimsPass());
+  // The resolution of `dim` ops tends to create identical ops. CSE them.
+  pm.addNestedPass<func::FuncOp>(createCSEPass());
+
+  // Finish the type conversion from `torch` types to the types of the
+  // TOSA backend contract.
+  pm.addPass(TorchConversion::createFuncBackendTypeConversionPass());
+  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+  pm.addNestedPass<func::FuncOp>(
+      TorchConversion::createFinalizingBackendTypeConversionPass());
+
+  // Verify that we have lowered to ops that are supported by union of TOSA and
+  // Linalg-on-tensors backend. This fails compilation (signalPassFailure) if
+  // the IR is not in the correct form.
+  pm.addPass(TorchConversion::createVerifyTosaLinalgBackendContractPass());
+}
 #endif
 
 #ifdef TORCH_MLIR_ENABLE_STABLEHLO
diff --git a/lib/Dialect/TorchConversion/Transforms/VerifyTosaLinalgBackendContract.cpp b/lib/Dialect/TorchConversion/Transforms/VerifyTosaLinalgBackendContract.cpp
diff --git a/projects/pt1/e2e_testing/main.py b/projects/pt1/e2e_testing/main.py
diff --git a/python/torch_mlir/compiler_utils.py b/python/torch_mlir/compiler_utils.py
diff --git a/test/Conversion/TorchToTosa/torch-backend-to-tosa-backend-pipeline.mlir b/test/Conversion/TorchToTosa/torch-backend-to-tosa-backend-pipeline.mlir
diff --git a/test/Conversion/TorchToTosa/torch-backend-to-tosa-linalg-backend-pipeline.mlir b/test/Conversion/TorchToTosa/torch-backend-to-tosa-linalg-backend-pipeline.mlir
diff --git a/utils/bazel/torch-mlir-overlay/BUILD.bazel b/utils/bazel/torch-mlir-overlay/BUILD.bazel
