From 5243638e3366ed0eba196f6382c0056ca70d94c7 Mon Sep 17 00:00:00 2001
From: Kristof Denolf <kristof@xilinx.com>
Date: Tue, 3 May 2022 12:57:33 -0700
Subject: [PATCH] add no decompose option

---
 include/torch-mlir/Dialect/Torch/Transforms/Passes.h | 5 +++++
 lib/Dialect/Torch/Transforms/Passes.cpp              | 4 +++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/include/torch-mlir/Dialect/Torch/Transforms/Passes.h b/include/torch-mlir/Dialect/Torch/Transforms/Passes.h
index 0f86364d0..790b135c9 100644
--- a/include/torch-mlir/Dialect/Torch/Transforms/Passes.h
+++ b/include/torch-mlir/Dialect/Torch/Transforms/Passes.h
@@ -32,6 +32,11 @@ struct TorchLoweringPipelineOptions
   // If this option is false, only do the bare minimum for correctness.
   Option<bool> optimize{*this, "optimize", llvm::cl::desc("Do optimizations."),
                         llvm::cl::init(true)};
+  
+  // If this option is false, decompose complex operations.
+  // If this option is true, skip decomposition of complex operations.
+  Option<bool> nodecompose{*this, "nodecompose", llvm::cl::desc("No decomposition."),
+                        llvm::cl::init(false)};                      
 };
 
 /// Creates a pipeline that lowers the object graph IR that is produced by
diff --git a/lib/Dialect/Torch/Transforms/Passes.cpp b/lib/Dialect/Torch/Transforms/Passes.cpp
index 94644c317..1f6b5176a 100644
--- a/lib/Dialect/Torch/Transforms/Passes.cpp
+++ b/lib/Dialect/Torch/Transforms/Passes.cpp
@@ -154,7 +154,9 @@ void mlir::torch::Torch::createTorchFunctionToTorchBackendPipeline(
     // only-used-in-training operations on `torch.global_slot`'s.
     pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   }
-  pm.addNestedPass<func::FuncOp>(Torch::createDecomposeComplexOpsPass());
+  
+  if (!options.nodecompose)
+    pm.addNestedPass<func::FuncOp>(Torch::createDecomposeComplexOpsPass());
 
   // TODO: VerifyTorchBackendContractPass.
 }