Merge branch 'main' into ferdinand.update_llvm_april_2024

Author: hamptonm1

src/Builder/FrontendDialectTransformer.cpp

@@ -51,6 +51,8 @@ SUPPRESS_WARNINGS_POP
 #include <array>
 #include <fstream>
 #include <map>
+#include <sstream>
+#include <string>
 #include <unordered_map>
 #include <vector>
 
@@ -450,6 +452,27 @@ class FrontendGenImpl {
     return attributes;
   }
 
+  // Generate a string vector from the dimParams option string
+  void getInputDimParamsMapFromOption(std::string optionStr,
+      std::map<int, std::string> &paramStrMap,
+      std::string &paramStrForAllArgs) {
+    std::stringstream paramStrStream(optionStr);
+    std::string dimParamStr;
+    while (std::getline(paramStrStream, dimParamStr, '|')) {
+      size_t pos = dimParamStr.find(':');
+      assert((pos > 0) && "invalid dimParams option string");
+      int idx = stoi(dimParamStr.substr(0, pos));
+      dimParamStr = dimParamStr.substr(pos + 1);
+      std::replace(dimParamStr.begin(), dimParamStr.end(), '=', ':');
+      if (idx < 0) // set all arguments
+        paramStrForAllArgs = dimParamStr;
+      else {
+        paramStrMap[idx] = dimParamStr;
+      }
+    }
+    return;
+  }
+
   /*!
    * An alternative graph importing procedure for importing ONNX subgraphs.
    * ONNX subgraphs, unlike the main computation graph, are imported as regions
@@ -490,6 +513,10 @@ class FrontendGenImpl {
     // See https://github.com/onnx/onnx/blob/main/docs/IR.md for more
     // information about dim_param.
     llvm::SmallVector<std::string, 4> inputDimParams, outputDimParams;
+    std::map<int, std::string> inputDimParamsFromOption;
+    std::string inputDimParamsFromOptionForAllArgs;
+    getInputDimParamsMapFromOption(options_.dimParams, inputDimParamsFromOption,
+        inputDimParamsFromOptionForAllArgs);
 
     // Import the input tensor types that are not constant and not initialized.
     int inputIndex = 0;
@@ -500,7 +527,16 @@ class FrontendGenImpl {
         std::string dimParams = "";
         Type argTy = ImportType(input.type(), &dimParams);
         argTy = modelInputShaper_.reshape(inputIndex, argTy);
-        if (!dimParams.empty())
+        // For each input tensor, use either all dimensions by the compiler
+        // option OR all dimensions in the original onnx model. Dimensions
+        // from the option and the model in a single input tensor are not
+        // merged.
+        if (inputDimParamsFromOption.find(inputIndex) !=
+            inputDimParamsFromOption.end())
+          inputDimParams.emplace_back(inputDimParamsFromOption[inputIndex]);
+        else if (!inputDimParamsFromOptionForAllArgs.empty())
+          inputDimParams.emplace_back(inputDimParamsFromOptionForAllArgs);
+        else if (!dimParams.empty())
           inputDimParams.emplace_back(dimParams);
 
         argTypes.emplace_back(argTy);

src/Builder/FrontendDialectTransformer.hpp

@@ -55,6 +55,18 @@ struct ImportOptions {
   //   - (arg0: tensor<3x4x5xf32>, arg1: tensor<10x5xf32>)
   //
   std::string shapeInformation = "";
+  // Custom onnx.dim_params attributes for the graph inputs for specifying
+  // relationship among their dynamic dimensions.
+  // Its format is 'input_id:dim_id=sym,dim_id=sym,...|input_id:
+  // dim_id=sym,dim_id=sym,...|input_id...'
+  // E.g. An ONNX model has two dynamic inputs
+  //   - (arg0: tensor<?x5xf32>, arg1: tensor<?x5xf32>)
+  // If we want to specify that the first unknown dimension of arg0 and the
+  // first unknown dimension of arg1 are the same, we can assign the two
+  // dimensions to the same symbol "batch" as follows.
+  //   - dimParams = '0:0=batch|1:0=batch'
+  //
+  std::string dimParams = "";
   // Directory to look for external data if any tensor has external
   // data location. If empty then external data is disabled.
   std::string externalDataDir = "";

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
@@ -51,6 +52,7 @@ bool preserveMLIR;                                     // onnx-mlir only
 bool useOnnxModelTypes;                                // onnx-mlir only
 int repeatOnnxTransform;                               // onnx-mlir only
 std::string shapeInformation;                          // onnx-mlir only
+std::string dimParams;                                 // onnx-mlir only
 ModelSize modelSize;                                   // onnx-mlir only
 bool storeConstantsToFile;                             // onnx-mlir only
 float constantsToFileTotalThreshold;                   // onnx-mlir only
@@ -201,6 +203,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),
@@ -282,6 +291,23 @@ static llvm::cl::opt<std::string, true> shapeInformationOpt("shapeInformation",
     llvm::cl::value_desc("value"), llvm::cl::location(shapeInformation),
     llvm::cl::cat(OnnxMlirOptions));
 
+static llvm::cl::opt<std::string, true> dimParamsOpt("dimParams",
+    llvm::cl::desc(
+        "Custom onnx.dim_params attributes for the inputs of the ONNX model for"
+        "specifying relationship among dynamic dimensions of the inputs.\n"
+        "\"value\" is in the format of "
+        "\"INPUT_ID1:D1=S1,D2=S2,...,Dn=Sn|INPUT_ID2:D1=T1,D2=T2,...Dn=Tn|"
+        "...\" where \"INPUT_ID1, INPUT_ID2, ...\" are input indices "
+        "(starting from 0 or being -1 for all input indices), and\n"
+        "\"S1, S2, ...\" and \"T2, T2, ...\" are symbols to specify that same "
+        "symbols have the same value. "
+        "All dimensions of onnx.dim_params for a specified input index in "
+        "the original onnx model are cleared and repalced by this option. "
+        "onnx.dim_params for other input indices in the original onnx model "
+        "are not cleared"),
+    llvm::cl::value_desc("value"), llvm::cl::location(dimParams),
+    llvm::cl::cat(OnnxMlirOptions));
+
 // Default value is defined by the OnnxMlirEnvOptionName constant string
 // variable, but the default setting mechanism here cannot be used here as we
 // need to evaluate this value prior to the compiler options being set. Proper

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
@@ -94,6 +95,7 @@ extern bool preserveMLIR;                                     // onnx-mlir only
 extern bool useOnnxModelTypes;                                // onnx-mlir only
 extern int repeatOnnxTransform;                               // onnx-mlir only
 extern std::string shapeInformation;                          // onnx-mlir only
+extern std::string dimParams;                                 // onnx-mlir only
 extern ModelSize modelSize;                                   // onnx-mlir only
 extern bool storeConstantsToFile;                             // onnx-mlir only
 extern float constantsToFileTotalThreshold;                   // onnx-mlir only

src/Compiler/CompilerUtils.cpp

@@ -636,6 +636,7 @@ int processInputFile(StringRef inputFilename, mlir::MLIRContext &context,
     options.useOnnxModelTypes = useOnnxModelTypes;
     options.invokeOnnxVersionConverter = invokeOnnxVersionConverter;
     options.shapeInformation = shapeInformation;
+    options.dimParams = dimParams;
     options.allowSorting = allowSorting;
     options.externalDataDir = dirName(inputFilename);
     options.functionsToDecompose.insert(options.functionsToDecompose.end(),

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.