Merge pull request #324 from Xilinx/matthias.onnx_options

Use pass options instead of global variables to drive Decompose

Author: mgehre-amd

src/Compiler/CompilerPasses.cpp

@@ -67,8 +67,7 @@ void configurePasses() {
 }
 
 void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
-    bool donotScrubDisposableElementsAttr,
-    bool enableQuarkQuantizedLegalization) {
+    bool donotScrubDisposableElementsAttr, OnnxToMlirOptions opts) {
   // This is a transition from previous static passes to full dynamic passes
   // Static passes are kept and the dynamic pass is added as IF-THEN
   // with the static iteration.
@@ -85,12 +84,14 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
         std::make_unique<DisposableGarbageCollector>(pm.getContext()));
 
   // Decompose first. Eliminates some unsupported ops without shape inference.
-  pm.addNestedPass<func::FuncOp>(onnx_mlir::createDecomposeONNXToONNXPass());
+  pm.addNestedPass<func::FuncOp>(onnx_mlir::createDecomposeONNXToONNXPass(
+      /*target=*/"", opts.enableConvTransposeDecompose,
+      opts.enableConvTranposeDecomposeToPhasedConv));
   if (!disableRecomposeOption)
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createRecomposeONNXToONNXPass());
   if (enableONNXHybridPass) {
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createONNXHybridTransformPass(
-        !disableRecomposeOption, enableQuarkQuantizedLegalization));
+        !disableRecomposeOption, opts.enableQuarkQuantizedLegalization));
     // Convolution Optimization for CPU: enable when there are no accelerators.
     if (targetCPU && enableConvOptPass) {
       pm.addNestedPass<func::FuncOp>(onnx_mlir::createConvOptONNXToONNXPass(
@@ -130,13 +131,13 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
 
   // Simplify shape-related ops.
   pm.addPass(onnx_mlir::createSimplifyShapeRelatedOpsPass(
-      enableQuarkQuantizedLegalization));
+      opts.enableQuarkQuantizedLegalization));
 
   // One more call to ONNX shape inference/canonicalization/... to update
   // shape if possible.
   if (enableONNXHybridPass) {
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createONNXHybridTransformPass(
-        !disableRecomposeOption, enableQuarkQuantizedLegalization));
+        !disableRecomposeOption, opts.enableQuarkQuantizedLegalization));
   } else {
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createShapeInferencePass());
     pm.addPass(mlir::createCanonicalizerPass());
@@ -335,10 +336,15 @@ void addPasses(mlir::OwningOpRef<ModuleOp> &module, mlir::PassManager &pm,
 
   // NOTE: FlexML sets the targetCPU flag to false, as we do not want to run
   //       the CPU specific transformations.
-  if (inputIRLevel <= ONNXLevel && emissionTarget >= EmitONNXIR)
+  if (inputIRLevel <= ONNXLevel && emissionTarget >= EmitONNXIR) {
+    OnnxToMlirOptions opts;
+    opts.enableQuarkQuantizedLegalization = enableQuarkQuantizedLegalization;
+    opts.enableConvTransposeDecompose = enableConvTransposeDecomposeOption;
+    opts.enableConvTranposeDecomposeToPhasedConv =
+        enableConvTranposeDecomposeToPhasedConv;
     addONNXToMLIRPasses(pm, /*target CPU*/ false,
-        /*donotScrubDisposableElementsAttr=*/false,
-        /*enableQuarkQuantizedLegalization=*/enableQuarkQuantizedLegalization);
+        /*donotScrubDisposableElementsAttr=*/false, opts);
+  }
 
   if (emissionTarget >= EmitMLIR) {
     if (inputIRLevel <= ONNXLevel)

src/Compiler/CompilerPasses.hpp

@@ -14,15 +14,25 @@
 
 #ifndef ONNX_MLIR_COMPILER_PASSES_H
 #define ONNX_MLIR_COMPILER_PASSES_H
+#include "mlir/IR/OwningOpRef.h"
 #include "mlir/Pass/PassManager.h"
+#include "onnx-mlir/Compiler/OMCompilerTypes.h"
+namespace mlir {
+class ModuleOp;
+}
 
 namespace onnx_mlir {
 // Configures passes up front based on command line options.
 void configurePasses();
 
+struct OnnxToMlirOptions {
+  bool enableQuarkQuantizedLegalization = false;
+  bool enableConvTransposeDecompose = false;
+  bool enableConvTranposeDecomposeToPhasedConv = false;
+};
+
 void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
-    bool donotScrubDisposableElementsAttr = false,
-    bool enableQuarkQuantizedLegalization = false);
+    bool donotScrubDisposableElementsAttr = false, OnnxToMlirOptions opts = {});
 void addONNXToKrnlPasses(mlir::PassManager &pm, int optLevel, bool enableCSE,
     std::string instrumentSignatureString, std::string ONNXOpsStatFilename);
 void addKrnlToAffinePasses(mlir::PassManager &pm);

src/Dialect/ONNX/Transforms/Decompose.cpp

@@ -33,7 +33,6 @@
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
-#include "src/Compiler/CompilerOptions.hpp"
 #include "src/Dialect/ONNX/DialectBuilder.hpp"
 #include "src/Dialect/ONNX/ElementsAttr/ElementsAttrHelper.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -49,6 +48,10 @@
 
 using namespace mlir;
 
+namespace {
+thread_local bool localEnableConvTransposeDecompose = false;
+thread_local bool localEnableConvTranposeDecomposeToPhasedConv = false;
+} // namespace
 namespace onnx_mlir {
 
 // Create an DenseElementsAttr of ArrayAttr.
@@ -659,7 +662,7 @@ Value replaceSequenceAt(
 }
 
 bool shouldDecomposeConvTransposeOp(Value convTransposeResult) {
-  if (!onnx_mlir::enableConvTransposeDecomposeOption) {
+  if (!localEnableConvTransposeDecompose) {
     // Disable the ONNXConvTransposeOp decomposition patterns.
     return false;
   }
@@ -714,7 +717,7 @@ bool hasNoActivationConsumer(Value convTransposeResult) {
 bool ShouldDecomposeConvTransposeOpToPhasedConvs(Value convTransposeResult,
     ArrayAttr kernelShapeAttr, ArrayAttr padsShapeAttr,
     ArrayAttr stridesShapeAttr, ArrayAttr outputShapeAttr) {
-  if (!onnx_mlir::enableConvTranposeDecomposeToPhasedConv) {
+  if (!localEnableConvTranposeDecomposeToPhasedConv) {
     // Disable the ONNXConvTransposeOp to Conv decomposition patterns.
     return false;
   }
@@ -2738,11 +2741,23 @@ struct DecomposeONNXToONNXPass
     : public PassWrapper<DecomposeONNXToONNXPass, OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DecomposeONNXToONNXPass)
 
-  DecomposeONNXToONNXPass(const std::string &target) { this->target = target; }
+  DecomposeONNXToONNXPass(const std::string &target,
+      bool enableConvTransposeDecompose = false,
+      bool enableConvTranposeDecomposeToPhasedConv = false) {
+    this->target = target;
+    this->enableConvTransposeDecompose = enableConvTransposeDecompose;
+    this->enableConvTranposeDecomposeToPhasedConv =
+        enableConvTranposeDecomposeToPhasedConv;
+  }
+
   DecomposeONNXToONNXPass(const DecomposeONNXToONNXPass &pass)
       : mlir::PassWrapper<DecomposeONNXToONNXPass,
             OperationPass<func::FuncOp>>() {
     this->target = pass.target.getValue();
+    this->enableConvTransposeDecompose =
+        pass.enableConvTransposeDecompose.getValue();
+    this->enableConvTranposeDecomposeToPhasedConv =
+        pass.enableConvTranposeDecomposeToPhasedConv.getValue();
   }
 
   StringRef getArgument() const override { return "decompose-onnx"; }
@@ -2755,6 +2770,16 @@ struct DecomposeONNXToONNXPass
   Option<std::string> target{*this, "target",
       llvm::cl::desc("Target Dialect to decompose into"), ::llvm::cl::init("")};
 
+  Option<bool> enableConvTransposeDecompose{*this, "enable-convtranspose",
+      llvm::cl::desc("Enable decomposition of ConvTranspose"),
+      ::llvm::cl::init(false)};
+
+  Option<bool> enableConvTranposeDecomposeToPhasedConv{*this,
+      "enable-convtranspose-phased",
+      llvm::cl::desc("Enable decomposition of ONNX ConvTranspose operator to 4 "
+                     "phased Conv"),
+      ::llvm::cl::init(false)};
+
   void runOnOperation() final;
 
   typedef PassWrapper<DecomposeONNXToONNXPass, OperationPass<func::FuncOp>>
@@ -2774,7 +2799,17 @@ void DecomposeONNXToONNXPass::runOnOperation() {
   }
 #endif
 
-  if (failed(applyPatternsAndFoldGreedily(function, std::move(patterns))))
+  // Set thread locals to affect native functions called by .td patterns.
+  localEnableConvTransposeDecompose = enableConvTransposeDecompose;
+  localEnableConvTranposeDecomposeToPhasedConv =
+      enableConvTranposeDecomposeToPhasedConv;
+
+  auto status = applyPatternsAndFoldGreedily(function, std::move(patterns));
+
+  localEnableConvTransposeDecompose = false;
+  localEnableConvTranposeDecomposeToPhasedConv = false;
+
+  if (failed(status))
     signalPassFailure();
 }
 
@@ -2809,6 +2844,8 @@ void onnx_mlir::getDecomposeONNXToONNXPatterns(
  * Create a DecomposeONNX pass.
  */
 std::unique_ptr<mlir::Pass> onnx_mlir::createDecomposeONNXToONNXPass(
-    const std::string &target) {
-  return std::make_unique<DecomposeONNXToONNXPass>(target);
-}
+    const std::string &target, bool enableConvTransposeDecompose,
+    bool enableConvTranposeDecomposeToPhasedConv) {
+  return std::make_unique<DecomposeONNXToONNXPass>(target,
+      enableConvTransposeDecompose, enableConvTranposeDecomposeToPhasedConv);
+}

src/Pass/Passes.hpp

@@ -38,7 +38,8 @@ std::unique_ptr<mlir::Pass> createONNXOpTransformPass(int threshold,
 
 /// Pass for rewriting inside frontend dialect.
 std::unique_ptr<mlir::Pass> createDecomposeONNXToONNXPass(
-    const std::string &target = "");
+    const std::string &target = "", bool enableConvTransposeDecompose = false,
+    bool enableConvTranposeDecomposeToPhasedConv = false);
 std::unique_ptr<mlir::Pass> createRecomposeONNXToONNXPass(
     const std::string &target = "");
 
@@ -132,4 +133,4 @@ std::unique_ptr<mlir::Pass> createConvertKrnlToLLVMPass(bool verifyInputTensors,
 /// Pass for lowering Onnx ops to TOSA dialect
 std::unique_ptr<mlir::Pass> createConvertONNXToTOSAPass();
 } // namespace onnx_mlir
-#endif
+#endif

test/mlir/onnx/onnx_decompose_convtranspose.mlir

@@ -1,4 +1,4 @@
-// RUN: onnx-mlir-opt --shape-inference --decompose-onnx --enable-convtranspose-decompose %s -split-input-file | FileCheck %s
+// RUN: onnx-mlir-opt --shape-inference --decompose-onnx=enable-convtranspose %s -split-input-file | FileCheck %s
 // RUN: onnx-mlir-opt --shape-inference --decompose-onnx %s -split-input-file | FileCheck %s --check-prefix=DISABLED
 
 // -----

test/mlir/onnx/onnx_decompose_convtranspose_phased_conv.mlir

@@ -1,4 +1,4 @@
-// RUN: onnx-mlir-opt --shape-inference --decompose-onnx --enable-convtranspose-decompose-phased-conv %s -split-input-file | FileCheck %s
+// RUN: onnx-mlir-opt --shape-inference --decompose-onnx=enable-convtranspose-phased %s -split-input-file | FileCheck %s
 
 func.func @test_convtrans_4phase_kernel_shape_66(%arg0: tensor<1x512x10x16xf32>, %arg1: tensor<512x256x6x6xf32>) -> tensor<1x256x20x32xf32> {    
   %0 = "onnx.Constant" () { value= dense<0.02> : tensor<256xf32>} : ()-> tensor<256xf32>