Add options to specify directory with ONNX initializer files

Xilinx · Feb 11, 2025 · 17fb894 · 17fb894
1 parent f4d1ac4
commit 17fb894
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 1 deletion.
diff --git a/src/Compiler/CompilerOptions.cpp b/src/Compiler/CompilerOptions.cpp
@@ -59,6 +59,7 @@ int repeatOnnxTransform;                               // onnx-mlir only
 std::string shapeInformation;                          // onnx-mlir only
 std::string dimParams;                                 // onnx-mlir only
 ModelSize modelSize;                                   // onnx-mlir only
+std::string externalDataDir;                           // onnx-mlir only
 bool storeConstantsToFile;                             // onnx-mlir only
 float constantsToFileTotalThreshold;                   // onnx-mlir only
 float constantsToFileSingleThreshold;                  // onnx-mlir only
@@ -392,6 +393,16 @@ static llvm::cl::opt<ModelSize, true> modelSizeOpt("modelSize",
     llvm::cl::init(small), llvm::cl::cat(OnnxMlirOptions),
     llvm::cl::ValueRequired);
 
+static llvm::cl::opt<std::string, true> externalDataDirOpt("external-data-dir",
+    llvm::cl::desc(
+        "ONNX constant initializers can be stored in a separate file. "
+        "The filename is stored in the ONNX model without a path. By default "
+        "the path is the same folder as the ONNX model. This allows that "
+        "default to be overridden.\n"
+        "Default is empty (use default)."),
+    llvm::cl::location(externalDataDir), llvm::cl::init(""),
+    llvm::cl::cat(OnnxMlirOptions));
+
 static llvm::cl::opt<bool, true> storeConstantsToFileOpt(
     "store-constants-to-file",
     llvm::cl::desc(

diff --git a/src/Compiler/CompilerOptions.hpp b/src/Compiler/CompilerOptions.hpp
@@ -104,6 +104,7 @@ 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 std::string externalDataDir;                           // onnx-mlir only
 extern bool storeConstantsToFile;                             // onnx-mlir only
 extern float constantsToFileTotalThreshold;                   // onnx-mlir only
 extern float constantsToFileSingleThreshold;                  // onnx-mlir only

diff --git a/src/Compiler/CompilerUtils.cpp b/src/Compiler/CompilerUtils.cpp
@@ -725,7 +725,8 @@ int processInputFile(StringRef inputFilename, mlir::MLIRContext &context,
     options.shapeInformation = shapeInformation;
     options.dimParams = dimParams;
     options.allowSorting = allowSorting;
-    options.externalDataDir = dirName(inputFilename);
+    options.externalDataDir =
+        externalDataDir.empty() ? dirName(inputFilename) : externalDataDir;
     options.functionsToDecompose.insert(options.functionsToDecompose.end(),
         functionsToDecompose.begin(), functionsToDecompose.end());
     return ImportFrontendModelFile(