Add an option to disable op fusion in ONNXToKrnl pass (#2807)

* parallel for

Signed-off-by: Chen Tong <[email protected]>

* run through

Signed-off-by: Tong Chen <[email protected]>

* flag

Signed-off-by: Tong Chen <[email protected]>

* format

Signed-off-by: chentong319 <[email protected]>

* Revert "format"

This reverts commit ebd27e2.

* Revert "flag"

This reverts commit 432e4be.

* fix merge error

Signed-off-by: chentong319 <[email protected]>

* add function

Signed-off-by: chentong319 <[email protected]>

* format

Signed-off-by: chentong319 <[email protected]>

* cmake

Signed-off-by: chentong319 <[email protected]>

---------

Signed-off-by: Chen Tong <[email protected]>
Signed-off-by: Tong Chen <[email protected]>
Signed-off-by: Tong Chen <[email protected]>
Signed-off-by: chentong319 <[email protected]>

Author: chentong319

src/Compiler/CompilerOptions.cpp

@@ -41,6 +41,7 @@ std::vector<std::string> onnxConstPropDisablePatterns; // common for both
 bool enableONNXHybridPass;                             // common for both
 std::vector<std::string> functionsToDecompose;         // common for both
 std::string opsForCall;                                // common for both
+bool disableKrnlOpFusion;                              // common for both
 EmissionTargetType emissionTarget;                     // onnx-mlir only
 bool invokeOnnxVersionConverter;                       // onnx-mlir only
 bool preserveLocations;                                // onnx-mlir only
@@ -201,6 +202,13 @@ static llvm::cl::list<std::string, std::vector<std::string>>
         llvm::cl::location(functionsToDecompose),
         llvm::cl::cat(OnnxMlirCommonOptions));
 
+static llvm::cl::opt<bool, true> disableKrnlOpFusionOpt(
+    "disable-krnl-op-fusion",
+    llvm::cl::desc("disable op fusion in onnx-to-krnl pass (default=false)\n"
+                   "Set to 'true' if you want to disable fusion."),
+    llvm::cl::location(disableKrnlOpFusion), llvm::cl::init(false),
+    llvm::cl::cat(OnnxMlirCommonOptions));
+
 static llvm::cl::opt<bool, true> disableRecomposeOptionOpt("disable-recompose",
     llvm::cl::desc("Disable recomposition of ONNX operations."),
     llvm::cl::location(disableRecomposeOption), llvm::cl::init(false),

src/Compiler/CompilerOptions.hpp

@@ -84,6 +84,7 @@ extern std::vector<std::string> onnxConstPropDisablePatterns; // common for both
 extern bool enableONNXHybridPass;                             // common for both
 extern std::vector<std::string> functionsToDecompose;         // common for both
 extern std::string opsForCall;                                // common for both
+extern bool disableKrnlOpFusion;                              // common for both
 extern EmissionTargetType emissionTarget;                     // onnx-mlir only
 extern bool invokeOnnxVersionConverter;                       // onnx-mlir only
 extern bool preserveLocations;                                // onnx-mlir only

src/Conversion/ONNXToKrnl/CMakeLists.txt

@@ -80,6 +80,7 @@ add_onnx_mlir_library(OMONNXToKrnl
 
   LINK_LIBS PUBLIC
   OMAccelerator
+  OMCompilerOptions
   OMONNXConversionCommon
   OMONNXOps
   OMSupport

src/Conversion/ONNXToKrnl/Math/Elementwise.cpp

@@ -17,6 +17,7 @@
 
 #include "llvm/Support/Debug.h"
 
+#include "src/Compiler/CompilerOptions.hpp"
 #include "src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp"
 #include "src/Dialect/Krnl/DialectBuilder.hpp"
 #include "src/Dialect/ONNX/ONNXOps/ShapeHelper.hpp"
@@ -1884,6 +1885,9 @@ bool OpFusionHelper::areInputsValidForFusion(
 // A successor op (user) is fusible if it is the only user, it is in the
 // fusible elementwise op list, and its inputs are valid for fusion.
 void OpFusionHelper::findFusibleOps() {
+  // Direct return if fusion is disabled
+  if (disableKrnlOpFusion)
+    return;
   Operation *defOp = rootOp;
   while (defOp->hasOneUse()) {
     // the possible ONNX Ops.