Add -march option to onnx-mlir (llvm#1107)

Signed-off-by: Ettore Tiotto <etiotto@ca.ibm.com>

include/OnnxMlirCompiler.h

@@ -43,12 +43,14 @@ namespace onnx_mlir {
  *  @param outputBaseName File name without extension to write output
  *  @param emissionTarget Target format to compile to
  *  @param mcpu Optional target CPU
+ *  @param march Optional target architecture
  *  @param mtriple Optional target architecture
  *  @return 0 on success or non-zero error code on failure
  */
 ONNX_MLIR_EXPORT int omCompileFromFile(const char *inputFilename,
     const char *outputBaseName, EmissionTargetType emissionTarget,
-    const char *mcpu, const char *mtriple, const char **errorMessage);
+    const char *mcpu, const char *march, const char *mtriple,
+    const char **errorMessage);
 
 /*!
  *  Compile an onnx model from an ONNX protobuf array
@@ -57,12 +59,13 @@ ONNX_MLIR_EXPORT int omCompileFromFile(const char *inputFilename,
  *  @param outputBaseName File name without extension to write output
  *  @param emissionTarget Target format to compile to
  *  @param mcpu Optional target CPU
+ *  @param march Optional target architecture
  *  @param mtriple Optional compile target triple
  *  @return 0 on success or non-zero error code on failure
  */
 ONNX_MLIR_EXPORT int omCompileFromArray(const void *inputBuffer, int bufferSize,
     const char *outputBaseName, EmissionTargetType emissionTarget,
-    const char *mcpu, const char *mtriple);
+    const char *mcpu, const char *march, const char *mtriple);
 
 } // namespace onnx_mlir
 

src/Compiler/CompilerUtils.cpp

@@ -95,14 +95,19 @@ static llvm::cl::opt<string> shapeInformation("shapeInformation",
     llvm::cl::value_desc("value"), llvm::cl::cat(OnnxMlirOptions));
 
 static llvm::cl::opt<std::string> mtriple("mtriple",
-    llvm::cl::desc("Target architecture"),
+    llvm::cl::desc("Override target triple for module"),
     llvm::cl::value_desc("LLVM target triple>"), llvm::cl::cat(OnnxMlirOptions),
     llvm::cl::ValueRequired);
 
 static llvm::cl::opt<std::string> mcpu("mcpu", llvm::cl::desc("Target cpu"),
     llvm::cl::value_desc("Target a specific CPU type"),
     llvm::cl::cat(OnnxMlirOptions), llvm::cl::ValueRequired);
 
+static llvm::cl::opt<std::string> march("march",
+    llvm::cl::desc("Target architecture to generate code for"),
+    llvm::cl::value_desc("Target a specific architecture type"),
+    llvm::cl::cat(OnnxMlirOptions), llvm::cl::ValueRequired);
+
 enum OptLevel { O0, O1, O2, O3 };
 static llvm::cl::opt<OptLevel> OptimizationLevel(
     llvm::cl::desc("Optimization levels:"),
@@ -300,6 +305,7 @@ struct Command {
 } // namespace
 
 void setTargetCPU(const std::string &cpu) { mcpu = cpu; }
+void setTargetArch(const std::string &arch) { march = arch; }
 void setTargetTriple(const std::string &triple) { mtriple = triple; }
 
 static void LoadMLIR(string inputFilename, mlir::MLIRContext &context,
@@ -331,6 +337,13 @@ static std::string getTargetCpuOption() {
   return targetOptions;
 }
 
+static std::string getTargetArchOption() {
+  string targetOptions = "";
+  if (march != "")
+    targetOptions += "--march=" + march;
+  return targetOptions;
+}
+
 static std::string getTargetTripleOption() {
   string targetOptions = "";
   // Command cannot tolerate extra spaces. Add only when needed.
@@ -390,6 +403,7 @@ static void genLLVMBitcode(const mlir::OwningModuleRef &module,
   Command optBitcode(/*exePath=*/!optPath.empty() ? optPath : kOptPath);
   optBitcode.appendStr(getOptimizationLevelOption())
       .appendStr(getTargetTripleOption())
+      .appendStr(getTargetArchOption())
       .appendStr(getTargetCpuOption())
       .appendList({"-o", optimizedBitcodePath})
       .appendStr(unoptimizedBitcodePath)
@@ -408,6 +422,9 @@ static std::string genModelObject(string bitcodePath, string outputBaseName) {
   string llcPath = getToolPath("llc");
   Command llvmToObj(/*exePath=*/!llcPath.empty() ? llcPath : kLlcPath);
   llvmToObj.appendStr(getOptimizationLevelOption())
+      .appendStr(getTargetTripleOption())
+      .appendStr(getTargetArchOption())
+      .appendStr(getTargetCpuOption())
       .appendStr("-filetype=obj")
       .appendStr("-relocation-model=pic")
       .appendList({"-o", modelObjPath})

src/Compiler/CompilerUtils.hpp

@@ -35,6 +35,7 @@ extern llvm::cl::OptionCategory OnnxMlirOptions;
 extern llvm::cl::opt<std::string> instrumentONNXOps;
 
 void setTargetCPU(const std::string &cpu);
+void setTargetArch(const std::string &arch);
 void setTargetTriple(const std::string &triple);
 
 std::string compileModuleToObject(

src/Compiler/OnnxMlirCompiler.cpp

@@ -5,10 +5,12 @@
 #include "OnnxMlirCompiler.h"
 #include "CompilerUtils.hpp"
 
-void setCompileContext(
-    mlir::MLIRContext &context, const char *mcpu, const char *mtriple) {
+void setCompileContext(mlir::MLIRContext &context, const char *mcpu,
+    const char *march, const char *mtriple) {
   if (mcpu)
     setTargetCPU(std::string(mcpu));
+  if (march)
+    setTargetArch(std::string(march));
   if (mtriple)
     setTargetTriple(std::string(mtriple));
 
@@ -19,11 +21,12 @@ extern "C" {
 namespace onnx_mlir {
 ONNX_MLIR_EXPORT int omCompileFromFile(const char *inputFilename,
     const char *outputBaseName, EmissionTargetType emissionTarget,
-    const char *mcpu, const char *mtriple, const char **errorMessage) {
+    const char *mcpu, const char *march, const char *mtriple,
+    const char **errorMessage) {
   mlir::OwningModuleRef module;
   mlir::MLIRContext context;
 
-  setCompileContext(context, mcpu, mtriple);
+  setCompileContext(context, mcpu, march, mtriple);
   std::string error_message;
   processInputFile(std::string(inputFilename), context, module, &error_message);
   if (errorMessage != NULL) {
@@ -35,11 +38,11 @@ ONNX_MLIR_EXPORT int omCompileFromFile(const char *inputFilename,
 
 ONNX_MLIR_EXPORT int omCompileFromArray(const void *inputBuffer, int bufferSize,
     const char *outputBaseName, EmissionTargetType emissionTarget,
-    const char *mcpu, const char *mtriple) {
+    const char *mcpu, const char *march, const char *mtriple) {
   mlir::OwningModuleRef module;
   mlir::MLIRContext context;
 
-  setCompileContext(context, mcpu, mtriple);
+  setCompileContext(context, mcpu, march, mtriple);
   processInputArray(inputBuffer, bufferSize, context, module);
   return compileModule(module, context, outputBaseName, emissionTarget);
 }

test/backend/common.py

@@ -105,6 +105,8 @@ def compile_model(model, emit):
     command_list = [TEST_DRIVER]
     if args.mcpu:
         command_list.append("--mcpu=" + args.mcpu)
+    if args.march:
+        command_list.append("--march=" + args.march)
     if args.mtriple:
         command_list.append("--mtriple=" + args.mtriple)
     if args.converter or name in variables.test_need_converter:

test/backend/variables.py

@@ -2,7 +2,7 @@
 
 ############################ variables.py #####################################
 #
-# Copyright 2021 The IBM Research Authors.
+# Copyright 2021-2022 The IBM Research Authors.
 #
 ################################################################################
 # Immutable global variables:
@@ -115,6 +115,12 @@ def get_args_from_env():
         default=os.getenv("TEST_MCPU", ""),
         help="target a specific cpu, passed to the compiler",
     )
+    parser.add_argument(
+        "--march",
+        type=str,
+        default=os.getenv("TEST_MARCH", ""),
+        help="target a specific architecture, passed to the compiler",
+    )
     parser.add_argument(
         "--converter",
         action="store_true",
@@ -139,6 +145,8 @@ def get_runtime_vars():
 
     if args.mcpu:
         print("  targeting cpu:", args.mcpu, file=sys.stderr)
+    if args.march:
+        print("  targeting arch:", args.march, file=sys.stderr)
     if args.mtriple:
         print("  targeting triple:", args.mtriple, file=sys.stderr)
 

test/compilerlib/CompilerLibTest.cpp

@@ -12,6 +12,7 @@ using namespace onnx_mlir;
 std::string testFileName;
 std::string outputBaseName;
 std::string mcpu;
+std::string march;
 std::string mtriple;
 bool compileFromFile = false;
 
@@ -29,6 +30,7 @@ bool compileFromFile = false;
 
 void readArg(const std::string &arg) {
   PARSE_ARG(mcpu, "--mcpu=");
+  PARSE_ARG(march, "--march=");
   PARSE_ARG(mtriple, "--mtriple=");
   PARSE_ARG(outputBaseName, "-o");
   PARSE_FLAG(compileFromFile, "--fromfile");
@@ -52,6 +54,7 @@ int main(int argc, char *argv[]) {
     const char *errorMessage = NULL;
     retVal = omCompileFromFile(testFileName.c_str(), outputBaseName.c_str(),
         onnx_mlir::EmitLib, mcpu.empty() ? nullptr : mcpu.c_str(),
+        march.empty() ? nullptr : march.c_str(),
         mtriple.empty() ? nullptr : mtriple.c_str(), &errorMessage);
     if (errorMessage != NULL) {
       std::cerr << errorMessage;
@@ -65,6 +68,7 @@ int main(int argc, char *argv[]) {
     retVal =
         omCompileFromArray(test.data(), test.size(), outputBaseName.c_str(),
             onnx_mlir::EmitLib, mcpu.empty() ? nullptr : mcpu.c_str(),
+            march.empty() ? nullptr : march.c_str(),
             mtriple.empty() ? nullptr : mtriple.c_str());
   }
   if (retVal != 0) {