Merge pull request #473 from Xilinx/bump_to_d2330df5

[AutoBump] Merge with fixes of d2330df (Oct 21) (89)

Author: mgehre-amd

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

@@ -84,6 +84,11 @@ void createTorchDynamoExportToTorchBackendPipeline(
 void createTorchFunctionToTorchBackendPipeline(
     OpPassManager &pm, const TorchLoweringPipelineOptions &options);
 
+/// Creates a pipeline that lowers the torch Onnx IR that is produced by
+/// Onnx import into the form expected by torch-verify-backend-contract.
+void createTorchOnnxToTorchBackendPipeline(
+    OpPassManager &pm, const TorchLoweringPipelineOptions &options);
+
 /// Creates a pipeline that simplifies the computations in the program.
 /// This pass does not do any global program restructuring -- it works entirely
 /// within a single semantic model of a `builtin.module` with

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -3715,6 +3715,12 @@ OpFoldResult AtenRemainderScalarOp::fold(FoldAdaptor adaptor) {
 //===----------------------------------------------------------------------===//
 
 OpFoldResult AtenAddIntOp::fold(FoldAdaptor adaptor) {
+  auto intLhs = dyn_cast_or_null<IntegerAttr>(adaptor.getA());
+  auto intRhs = dyn_cast_or_null<IntegerAttr>(adaptor.getB());
+  if (intRhs && intRhs.getValue().getSExtValue() == 0)
+    return getA();
+  if (intLhs && intLhs.getValue().getSExtValue() == 0)
+    return getB();
   return atenBinaryIntOperatorFoldHelper(
       adaptor.getOperands(), [](int64_t a, int64_t b) { return a + b; });
 }
@@ -3724,6 +3730,9 @@ OpFoldResult AtenAddIntOp::fold(FoldAdaptor adaptor) {
 //===----------------------------------------------------------------------===//
 
 OpFoldResult AtenSubIntOp::fold(FoldAdaptor adaptor) {
+  if (getA() == getB())
+    return IntegerAttr::get(
+        IntegerType::get(getContext(), 64, IntegerType::Signless), 0);
   return atenBinaryIntOperatorFoldHelper(
       adaptor.getOperands(), [](int64_t a, int64_t b) { return a - b; });
 }
@@ -4590,7 +4599,8 @@ OpFoldResult AtenItemOp::fold(FoldAdaptor adaptor) {
     if (auto intAttr = dyn_cast<IntegerAttr>(splat)) {
       return intAttr.getType().isUnsignedInteger()
                  ? getI64IntegerAttr(getContext(), intAttr.getUInt())
-                 : getI64IntegerAttr(getContext(), intAttr.getSInt());
+                 : getI64IntegerAttr(getContext(),
+                                     intAttr.getValue().getSExtValue());
     }
     if (auto floatAttr = dyn_cast<FloatAttr>(splat)) {
       return getF64FloatAttr(getContext(), floatAttr.getValueAsDouble());

lib/Dialect/Torch/Transforms/Passes.cpp

@@ -10,6 +10,7 @@
 #include "torch-mlir/Dialect/Torch/Transforms/Passes.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
+#include "torch-mlir/Conversion/TorchOnnxToTorch/Passes.h"
 
 void mlir::torch::registerTorchPasses() {
   mlir::torch::registerPasses();
@@ -25,6 +26,10 @@ void mlir::torch::registerTorchPasses() {
       "torch-function-to-torch-backend-pipeline",
       "Pipeline lowering a Torch function to Torch backend form.",
       mlir::torch::Torch::createTorchFunctionToTorchBackendPipeline);
+  mlir::PassPipelineRegistration<Torch::TorchLoweringPipelineOptions>(
+      "torch-onnx-to-torch-backend-pipeline",
+      "Pipeline lowering Torch Onnx IR to Torch backend form.",
+      mlir::torch::Torch::createTorchOnnxToTorchBackendPipeline);
   mlir::PassPipelineRegistration<Torch::TorchLoweringPipelineOptions>(
       "torch-simplification-pipeline",
       "Pipeline simplifying computations in the program.",
@@ -86,6 +91,37 @@ void mlir::torch::Torch::createTorchFunctionToTorchBackendPipeline(
       options.backendLegalOps, options.extraLibrary));
 }
 
+void mlir::torch::Torch::createTorchOnnxToTorchBackendPipeline(
+    OpPassManager &pm, const TorchLoweringPipelineOptions &options) {
+  pm.addNestedPass<func::FuncOp>(onnx_c::createTorchOnnxToTorchPass());
+  // The above pass just converts the torch onnx IR to torch, hence the given
+  // pipeline will make sure that the IR is transformed such that it satisfies
+  // the backend contract.
+  if (options.decompose) {
+    pm.addNestedPass<func::FuncOp>(
+        Torch::createDecomposeComplexOpsPass(options.backendLegalOps));
+    pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+  }
+  // TODO: Move the combination of two passes i.e., ScalarizeShapes and
+  // TorchShapeRefinementPipeline out of here and create an onnx shape
+  // refinement pipeline which runs iteratively over the IR.
+  createTorchShapeRefinementPipeline(pm, options);
+  // This pass scalarizes the tensor shape computations.
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::torch::Torch::createScalarizeShapesPass());
+  createTorchShapeRefinementPipeline(pm, options);
+  pm.addPass(Torch::createRefinePublicReturnPass());
+  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+  // The decompose pass is run again here since the scalarize shapes pass and
+  // shape refinement pipeline might create some ops for which decomposition
+  // exists.
+  if (options.decompose) {
+    pm.addNestedPass<func::FuncOp>(
+        Torch::createDecomposeComplexOpsPass(options.backendLegalOps));
+    pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+  }
+}
+
 // A simplification pipeline to establish the invariants of the backend
 // contract (see `satisfiedBackendContract` in `LowerToBackendContract`).
 //