add options to specify library (#2407)

* add options Signed-off-by: chentong319 <[email protected]> * format Signed-off-by: chentong319 <[email protected]> * change flag Signed-off-by: chentong319 <[email protected]> * new approach Signed-off-by: chentong319 <[email protected]> * test flag Signed-off-by: chentong319 <[email protected]> * debug Signed-off-by: chentong319 <[email protected]> * clear config Signed-off-by: chentong319 <[email protected]> * format Signed-off-by: chentong319 <[email protected]> --------- Signed-off-by: chentong319 <[email protected]>
onnx · Aug 8, 2023 · 2fac346 · 2fac346
1 parent ebac513
commit 2fac346
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 4 deletions.
diff --git a/src/Compiler/CompilerOptions.cpp b/src/Compiler/CompilerOptions.cpp
@@ -305,6 +305,22 @@ llvm::cl::opt<std::string> modelTag("tag",
     llvm::cl::value_desc("a string that matches regex ([0-9a-z_.-]+)"),
     llvm::cl::init(""), llvm::cl::cat(OnnxMlirOptions));
 
+llvm::cl::opt<bool> enableConvOptPass("enable-conv-opt-pass",
+    llvm::cl::desc("Enable the ConvOptPass. Default is true."),
+    llvm::cl::init(true), llvm::cl::cat(OnnxMlirOptions));
+
+llvm::cl::list<std::string> extraLibPaths("L",
+    llvm::cl::desc("Specify extra directories for libraries when compiling"
+                   "an onnx model. Will be add used as -L in the linkage step."
+                   "Each directory can be specified with one extra-lib-dirs"),
+    llvm::cl::Prefix, llvm::cl::cat(OnnxMlirOptions));
+
+llvm::cl::list<std::string> extraLibs("l",
+    llvm::cl::desc("Specify extra libraries when compiling an onnx model."
+                   "Will be add used as -l in the linkage step."
+                   "Each lib can be specified with one extra-libs"),
+    llvm::cl::Prefix, llvm::cl::cat(OnnxMlirOptions));
+
 llvm::cl::opt<ProfileIRs> profileIR("profile-ir",
     llvm::cl::desc("Profile operations in an IR"),
     llvm::cl::values(clEnumVal(None, "No profiling. Default value."),
@@ -667,6 +683,9 @@ int setCompilerOptions(const CompilerOptionList &list) {
   return CompilerSuccess;
 }
 
+// Clear the map for CompilerConfig. It is used for each invocation of compile
+void clearCompilerConfig() { CompilerConfigMap.clear(); }
+
 // Get the string vector associated with the specified key
 std::vector<std::string> getCompilerConfig(std::string k) {
   return CompilerConfigMap[k];

diff --git a/src/Compiler/CompilerOptions.hpp b/src/Compiler/CompilerOptions.hpp
@@ -95,6 +95,9 @@ extern llvm::cl::opt<bool> enableSimdDataLayout;
 extern llvm::cl::opt<bool> enableONNXHybridPass;
 extern llvm::cl::list<std::string> functionsToDecompose;
 extern llvm::cl::opt<std::string> modelTag;
+extern llvm::cl::opt<bool> enableConvOptPass;
+extern llvm::cl::list<std::string> extraLibPaths;
+extern llvm::cl::list<std::string> extraLibs;
 extern llvm::cl::opt<ProfileIRs> profileIR;
 
 // The customEnvFlags must be scanned before the normal options.
@@ -143,6 +146,7 @@ using CompilerOptionList =
     llvm::SmallVector<std::pair<onnx_mlir::OptionKind, std::string>, 4>;
 
 #define CCM_SHARED_LIB_DEPS "sharedLibDeps"
+#define CCM_SHARED_LIB_PATH_DEPS "sharedLibPathDeps"
 extern std::map<std::string, std::vector<std::string>> CompilerConfigMap;
 
 // Return 0 on success. These functions are not thread-safe and should be called
@@ -154,6 +158,7 @@ std::string getCompilerOption(const onnx_mlir::OptionKind kind);
 
 // The add and del functions are not thread-safe and should only be
 // called from one thread.
+void clearCompilerConfig();
 std::vector<std::string> getCompilerConfig(std::string k);
 void addCompilerConfig(std::string k, std::vector<std::string> v);
 void delCompilerConfig(std::string k, std::vector<std::string> v);

diff --git a/src/Compiler/CompilerPasses.cpp b/src/Compiler/CompilerPasses.cpp
@@ -71,7 +71,7 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU) {
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createShapeInferencePass());
   }
   // Convolution Optimization for CPU: enable when there are no accelerators.
-  if (targetCPU) {
+  if (targetCPU && enableConvOptPass) {
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createConvOptONNXToONNXPass(
         enableSimdDataLayout && !disableSimdOption));
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createShapeInferencePass());

diff --git a/src/Compiler/CompilerUtils.cpp b/src/Compiler/CompilerUtils.cpp
@@ -583,7 +583,8 @@ static int compileModuleToSharedLibrary(
       modelObjNameWithExt, !keepFiles(KeepFilesOfType::Object));
   libNameWithExt = getTargetFilename(outputNameNoExt, EmitLib);
   return genSharedLib(libNameWithExt, {}, {modelObjNameWithExt},
-      getCompilerConfig(CCM_SHARED_LIB_DEPS), {getLibraryPath()});
+      getCompilerConfig(CCM_SHARED_LIB_DEPS),
+      getCompilerConfig(CCM_SHARED_LIB_PATH_DEPS));
 }
 
 // Return 0 on success, error code on failure
@@ -626,7 +627,7 @@ static int compileModuleToJniJar(
   std::string modelSharedLibPath = getTargetFilename(jniLibBase, EmitLib);
   rc = genSharedLib(modelSharedLibPath, NOEXECSTACK,
       {modelObjNameWithExt, jniObjPath}, getCompilerConfig(CCM_SHARED_LIB_DEPS),
-      {getLibraryPath()});
+      getCompilerConfig(CCM_SHARED_LIB_PATH_DEPS));
   if (rc != CompilerSuccess)
     return rc;
   llvm::FileRemover modelSharedLibRemover(
@@ -759,6 +760,13 @@ static int emitOutputFiles(std::string outputNameNoExt,
   } break;
   case EmitLib: {
     addCompilerConfig(CCM_SHARED_LIB_DEPS, {"cruntime"});
+    addCompilerConfig(CCM_SHARED_LIB_PATH_DEPS, {getLibraryPath()});
+    // Add user specified libs and their path
+    // Multiple lib or directory can be specified with multiple options.
+    // For example, -lextra1, -lextra2, -Lpath1, -Lpath2
+    addCompilerConfig(CCM_SHARED_LIB_DEPS, extraLibs);
+    addCompilerConfig(CCM_SHARED_LIB_PATH_DEPS, extraLibPaths);
+
     std::string sharedLibNameWithExt;
     int rc = compileModuleToSharedLibrary(
         module, outputNameNoExt, sharedLibNameWithExt);
@@ -775,6 +783,13 @@ static int emitOutputFiles(std::string outputNameNoExt,
   } break;
   case EmitJNI: {
     addCompilerConfig(CCM_SHARED_LIB_DEPS, {"jniruntime", "cruntime"});
+    addCompilerConfig(CCM_SHARED_LIB_PATH_DEPS, {getLibraryPath()});
+    // Add user specified libs and their path
+    // Multiple lib or directory can be specified with multiple options.
+    // For example, -lextra1, -lextra2, -Lpath1, -Lpath2
+    addCompilerConfig(CCM_SHARED_LIB_DEPS, extraLibs);
+    addCompilerConfig(CCM_SHARED_LIB_PATH_DEPS, extraLibPaths);
+
     int rc = compileModuleToJniJar(module, outputNameNoExt);
     if (rc != CompilerSuccess)
       return rc;
@@ -941,6 +956,8 @@ int compileModule(mlir::OwningOpRef<ModuleOp> &module,
   if (rc != CompilerSuccess)
     return rc;
 
+  clearCompilerConfig();
+
   configurePasses();
 
   mlir::PassManager pm(

diff --git a/src/Transform/ONNX/ONNXOpTransformPass.cpp b/src/Transform/ONNX/ONNXOpTransformPass.cpp
@@ -17,6 +17,7 @@
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
 
+#include "src/Compiler/CompilerOptions.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 #include "src/Pass/Passes.hpp"
 
@@ -78,7 +79,7 @@ void ONNXOpTransformPass::runOnOperation() {
     dynamicPM.addNestedPass<func::FuncOp>(
         onnx_mlir::createShapeInferencePass());
     // Convolution Optimization currently only for CPU.
-    if (onnxOpTransformTargetCPU) {
+    if (onnxOpTransformTargetCPU && onnx_mlir::enableConvOptPass) {
       dynamicPM.addNestedPass<func::FuncOp>(
           onnx_mlir::createConvOptONNXToONNXPass(
               onnxOpTransformEnableSimdDataLayout));