[mlir][mlir-cpu-runner] Add support for `-mattr` and `-march` flags

[awarzynski]

Summary

This patch adds support for -mattr and -march in mlir-cpu-runner.
With this change, one should be able to consistently use mlir-cpu-runner
for MLIR's integration tests (instead of e.g. resorting to lli when some
additional flags are needed). This is demonstrated in
concatenate_dim_1.mlir.

In order to support the new flags, this patch makes sure that
MLIR's ExecutionEngine/JITRunner (that mlir-cpu-runner is built on top of):

• takes into account the new command line flags when creating TargetMachine,
• avoids recreating TargetMachine if one is already available,
• creates LLVM's DataLayout based on the previously configured TargetMachine.

This is necessary in order to make sure that the command line
configuration is propagated correctly to the backend code generator.

A few additional updates are made in order to facilitate this change,
including support for debug dumps from JITRunner.

[awarzynski]

Fix Toy examples builds

[awarzynski]

Fix verify-flags.mlir

[c-rhodes]

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

190

nit: layout

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

230

nit: drop braces on single-statement ifs? (and below)

246

I think the previous error message was fine

251–253

what happens to the error message from createDefaultTargetMachine here?

mlir/lib/ExecutionEngine/JitRunner.cpp

411

nit: unrelated change

mlir/test/Integration/Dialect/SparseTensor/CPU/concatenate_dim_1.mlir

15

this test has been removed?

mlir/test/mlir-cpu-runner/verify-flags.mlir

2

Please add a test for march as well

[awarzynski]

Thanks for the review @c-rhodes , sending updates shortly.

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

251–253

Good catch, this should be propagated.

mlir/test/Integration/Dialect/SparseTensor/CPU/concatenate_dim_1.mlir

15

this test has been removed?

Yes, @aartbik and I discussed this offline and feel that this approach makes more sense (and this patch enables that).

Note that the 3rd and 4th RUN lines are exercising similar codepaths when ENABLE_VLA is false. Since these tests are expensive to run (and there's quite a few of them), it makes a lot of sense to reduce duplication. The solution is to replace 2 RUN lines with 1 in which the vectorizer will:

• generate regular fixed-vector width code when ENABLE_VLA is Off
• generate scalable SVE code when ENABLE_VLA is On.

Does it make sense?

mlir/test/mlir-cpu-runner/verify-flags.mlir

2

Good point.

For this to work, (i.e. for mlir-cpu-runner to fail in a sane way when using -march=<invalid-val>), I need to fix how errors are handled in JitRunner.cpp (just a wee heads-up)

[awarzynski]

Address Cullen's comments

In order to add a test for -march, I had to fix how errors are handled in JitRunner.cpp. Other changes should be self-explenatory.

[c-rhodes]

LGTM, cheers

mlir/test/Integration/Dialect/SparseTensor/CPU/concatenate_dim_1.mlir

15

this test has been removed?

Yes, @aartbik and I discussed this offline and feel that this approach makes more sense (and this patch enables that).

Note that the 3rd and 4th RUN lines are exercising similar codepaths when ENABLE_VLA is false. Since these tests are expensive to run (and there's quite a few of them), it makes a lot of sense to reduce duplication. The solution is to replace 2 RUN lines with 1 in which the vectorizer will:

• generate regular fixed-vector width code when ENABLE_VLA is Off
• generate scalable SVE code when ENABLE_VLA is On.

Does it make sense?

It does, thanks for clarifying.

[mehdi_amini]

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

240

Can we update this code to use JITTargetMachineBuilder (and JITTargetMachineBuilder::detectHost) instead of doing it manually here?

272

I have some concerns that your change here is going against the direction hinted by this TODO, what is the future plan about this? (your update to the TODO does not carve a path out)

[awarzynski]

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

240

Good point, this is basically re-implementing detectHost. I will delete it altogether.

272

This is a bit tricky :)

A triple could be defined in 3 different places (perhaps more?):

1. through methods like JITTargetMachineBuilder::detectHost/lvm::sys::getHostCPUName (the default value),
2. through command line arguments, e.g. -march, -target, -triple (the user-specified value),
3. inside the input LLVM module (the "input" defined value or via MLIR conversion/translation).

Prior to this change, only Option 1. was supported. With this change, mlir-cpu-runner will also support Option 2. This FIXME refers to Option 3. I replaced it with a TODO, because FIXME suggests that something is broken, but IMO that's not the case. It just happens that Options 2. and 3. are not supported ATM. My goal in this patch is to add support for Option 2 while keeping the path open for Option 3..

I said that this is tricky, because with 3 sources of the same information, one has to decide what takes precedence and then make sure that that's what's actually implemented. Clang and Flang have similar problem and approach it by ignoring Option 3 (tested on MacOS):

$ cat file.ll
target triple = "some-triple-different-to-host-triple"

define void @foo() {
  ret void
}
$ clang  -S file.ll
warning: overriding the module target triple with arm64-apple-macosx13.0.0 [-Woverride-module]
1 warning generated.

Sorry, this is a bit longer than I expected. The point that I'm trying to make is that:

• This patch only adds support for Option 2,
• Option 3. remains a TODO.

Is this change going to make Option 3. harder? I don't think so, but we'd need a motivating example to better understand what might be needed. We had a somewhat similar issue in Flang [1], and that wasn't too hard to support that (though that was specifically about DataLayout).

[1] https://github.com/llvm/llvm-project/issues/57230

[awarzynski]

Delete createDefaultTargetMachine

[awarzynski]

Fix build failures in the Tay examples

[awarzynski]

The outstanding pre-commit failure is unrelated to this change:

********************
Failed Tests (1):
  Flang :: Driver/omp-frontend-forwarding.f90
 
Testing Time: 15.32s
  Unsupported      :   48
  Passed           : 2019
  Expectedly Failed:   15
  Failed           :    1

[mehdi_amini]

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

272

Is this change going to make Option 3. harder? I don't think so, but we'd need a motivating example to better understand what might be needed

Basically that was the main thing I wanted to understand, I'm not sure I totally have wrapped my head entirely about this, but I'm willing to trust your judgement on this :)

[awarzynski]

Btw, this patch refines how diagnostics are propagated, so you might see this:

$ ninja check-mlir
[35/36] Running the MLIR regression tests
Symbols not found: [ llvm_orc_registerEHFrameSectionWrapper ]

Testing Time: 5.63s
  Unsupported      :  190
  Passed           : 1813
  Expectedly Failed:    1

This comes from mlir-cpu-runner -host-supports-jit:

bin/mlir-cpu-runner --host-supports-jit
false
Symbols not found: [ llvm_orc_registerEHFrameSectionWrapper ]

It's a genuine diagnostic that wasn't reported before, but the error itself is bogus. I'm trying to fix it here: https://reviews.llvm.org/D146935 (however, that made me discover https://github.com/llvm/llvm-project/issues/61856).

mlir/examples/toy/Ch6/toyc.cpp

int dumpLLVMIR(mlir::ModuleOp module) {
   if (!llvmModule) {
     llvm::errs() << "Failed to emit LLVM IR\n";
     return -1;
   }
 
   // Initialize LLVM targets.
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
-  mlir::ExecutionEngine::setupTargetTriple(llvmModule.get());
+
+  // Configure the LLVM Module
+  auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
+  if (!tmBuilderOrError) {
+    llvm::errs() << "Could not create JITTargetMachineBuilder\n";
+    return -1;
+  }
+
+  auto tmOrError = tmBuilderOrError->createTargetMachine();
+  if (!tmOrError) {
+    llvm::errs() << "Could not create TargetMachine\n";
+    return -1;
+  }
+  mlir::ExecutionEngine::setupTargetTripleAndDataLayout(llvmModule.get(),
+                                                        tmOrError.get().get());
 
   /// Optionally run an optimization pipeline over the llvm module.
   auto optPipeline = mlir::makeOptimizingTransformer(
       /*optLevel=*/enableOpt ? 3 : 0, /*sizeLevel=*/0,
       /*targetMachine=*/nullptr);
   if (auto err = optPipeline(llvmModule.get())) {
     llvm::errs() << "Failed to optimize LLVM IR " << err << "\n";
     return -1;

mlir/examples/toy/Ch7/toyc.cpp

int dumpLLVMIR(mlir::ModuleOp module) {
   if (!llvmModule) {
     llvm::errs() << "Failed to emit LLVM IR\n";
     return -1;
   }
 
   // Initialize LLVM targets.
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
-  mlir::ExecutionEngine::setupTargetTriple(llvmModule.get());
+
+  // Create target machine and configure the LLVM Module
+  auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
+  if (!tmBuilderOrError) {
+    llvm::errs() << "Could not create JITTargetMachineBuilder\n";
+    return -1;
+  }
+
+  auto tmOrError = tmBuilderOrError->createTargetMachine();
+  if (!tmOrError) {
+    llvm::errs() << "Could not create TargetMachine\n";
+    return -1;
+  }
+  mlir::ExecutionEngine::setupTargetTripleAndDataLayout(llvmModule.get(),
+                                                        tmOrError.get().get());
 
   /// Optionally run an optimization pipeline over the llvm module.
   auto optPipeline = mlir::makeOptimizingTransformer(
       /*optLevel=*/enableOpt ? 3 : 0, /*sizeLevel=*/0,
       /*targetMachine=*/nullptr);
   if (auto err = optPipeline(llvmModule.get())) {
     llvm::errs() << "Failed to optimize LLVM IR " << err << "\n";
     return -1;

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

 /// arguments of the function, followed by a pointer to the result.  This allows
 /// the engine to provide the caller with a generic function pointer that can
 /// be used to invoke the JIT-compiled function.
 class ExecutionEngine {
 public:
   ExecutionEngine(bool enableObjectDump, bool enableGDBNotificationListener,
                   bool enablePerfNotificationListener);
 
-  /// Creates an execution engine for the given MLIR IR.
+  /// Creates an execution engine for the given MLIR IR. If TargetMachine is
+  /// not provided, default TM is created (i.e. ignoring any command line flags
+  /// that could affect the set-up).
   static llvm::Expected<std::unique_ptr<ExecutionEngine>>
-  create(Operation *op, const ExecutionEngineOptions &options = {});
+  create(Operation *op, const ExecutionEngineOptions &options = {},
+         std::unique_ptr<llvm::TargetMachine> tm = nullptr);
 
   /// Looks up a packed-argument function wrapping the function with the given
   /// name and returns a pointer to it. Propagates errors in case of failure.
   llvm::Expected<void (*)(void **)> lookupPacked(StringRef name) const;
 
   /// Looks up the original function with the given name and returns a
   /// pointer to it. This is not necesarily a packed function. Propagates
   /// errors in case of failure.
const std::string adapterName =
         std::string("_mlir_ciface_") + funcName.str();
     llvm::SmallVector<void *> argsArray;
     // Pack every arguments in an array of pointers. Delegate the packing to a
     // trait so that it can be overridden per argument type.
     (Argument<Args>::pack(argsArray, args), ...);
     return invokePacked(adapterName, argsArray);
   }
 
-  /// Set the target triple on the module. This is implicitly done when creating
-  /// the engine.
-  static bool setupTargetTriple(llvm::Module *llvmModule);
+  /// Set the target triple and the data layout for the input module based on
+  /// the input TargetMachine. This is implicitly done when creating the
+  /// engine.
+  static void setupTargetTripleAndDataLayout(llvm::Module *llvmModule,
+                                             llvm::TargetMachine *tm);
 
   /// Dump object code to output file `filename`.
   void dumpToObjectFile(StringRef filename);
 
   /// Register symbols with this ExecutionEngine.
   void registerSymbols(
       llvm::function_ref<llvm::orc::SymbolMap(llvm::orc::MangleAndInterner)>
           symbolMap);

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

 void ExecutionEngine::registerSymbols(
     llvm::function_ref<SymbolMap(MangleAndInterner)> symbolMap) {
   auto &mainJitDylib = jit->getMainJITDylib();
   cantFail(mainJitDylib.define(
       absoluteSymbols(symbolMap(llvm::orc::MangleAndInterner(
           mainJitDylib.getExecutionSession(), jit->getDataLayout())))));
 }
 
-// Setup LLVM target triple from the current machine.
-bool ExecutionEngine::setupTargetTriple(Module *llvmModule) {
-  // Setup the machine properties from the current architecture.
-  auto targetTriple = llvm::sys::getDefaultTargetTriple();
-  std::string errorMessage;
-  const auto *target =
-      llvm::TargetRegistry::lookupTarget(targetTriple, errorMessage);
-  if (!target) {
-    errs() << "NO target: " << errorMessage << "\n";
-    return true;
-  }
-
-  std::string cpu(llvm::sys::getHostCPUName());
-  llvm::SubtargetFeatures features;
-  llvm::StringMap<bool> hostFeatures;
-
-  if (llvm::sys::getHostCPUFeatures(hostFeatures))
-    for (const auto &[feature, isEnabled] : hostFeatures)
-      features.AddFeature(feature, isEnabled);
-
-  std::unique_ptr<llvm::TargetMachine> machine(target->createTargetMachine(
-      targetTriple, cpu, features.getString(), {}, {}));
-  if (!machine) {
-    errs() << "Unable to create target machine\n";
-    return true;
-  }
-  llvmModule->setDataLayout(machine->createDataLayout());
-  llvmModule->setTargetTriple(targetTriple);
-  return false;
+void ExecutionEngine::setupTargetTripleAndDataLayout(Module *llvmModule,
+                                                     llvm::TargetMachine *tm) {
+  llvmModule->setDataLayout(tm->createDataLayout());
+  llvmModule->setTargetTriple(tm->getTargetTriple().getTriple());
 }
 
 static std::string makePackedFunctionName(StringRef name) {
   return "_mlir_" + name.str();
 }
 
 // For each function in the LLVM module, define an interface function that wraps
 // all the arguments of the original function and all its results into an i8**
if (auto *listener = llvm::JITEventListener::createPerfJITEventListener())
       perfListener = listener;
     else if (auto *listener =
                  llvm::JITEventListener::createIntelJITEventListener())
       perfListener = listener;
   }
 }
 
 Expected<std::unique_ptr<ExecutionEngine>>
-ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options) {
+ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
+                        std::unique_ptr<llvm::TargetMachine> tm) {
   auto engine = std::make_unique<ExecutionEngine>(
       options.enableObjectDump, options.enableGDBNotificationListener,
       options.enablePerfNotificationListener);
 
   // Remember all entry-points if object dumping is enabled.
   if (options.enableObjectDump) {
     for (auto funcOp : m->getRegion(0).getOps<LLVM::LLVMFuncOp>()) {
       StringRef funcName = funcOp.getSymName();
       engine->functionNames.push_back(funcName.str());
     }
   }
 
   std::unique_ptr<llvm::LLVMContext> ctx(new llvm::LLVMContext);
   auto llvmModule = options.llvmModuleBuilder
                         ? options.llvmModuleBuilder(m, *ctx)
                         : translateModuleToLLVMIR(m, *ctx);
   if (!llvmModule)
     return makeStringError("could not convert to LLVM IR");
-  // FIXME: the triple should be passed to the translation or dialect conversion
-  // instead of this.  Currently, the LLVM module created above has no triple
-  // associated with it.
-  setupTargetTriple(llvmModule.get());
+
+  // If no valid TargetMachine was passed, create a default TM ignoring any
+  // input arguments from the user.
+  if (!tm) {
+    auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
+    if (!tmBuilderOrError)
+      return tmBuilderOrError.takeError();
+
+    auto tmOrError = tmBuilderOrError->createTargetMachine();
+    if (!tmOrError)
+      return tmOrError.takeError();
+    tm = std::move(tmOrError.get());
+  }
+
+  // TODO: Currently, the LLVM module created above has no triple associated
+  // with it. Instead, the triple is extracted from the TargetMachine, which is
+  // either based on the host defaults or command line arguments when specified
+  // (set-up by callers of this method). It could also be passed to the
+  // translation or dialect conversion instead of this.
+  setupTargetTripleAndDataLayout(llvmModule.get(), tm.get());
   packFunctionArguments(llvmModule.get());
 
   auto dataLayout = llvmModule->getDataLayout();
 
   // Callback to create the object layer with symbol resolution to current
   // process and dynamically linked libraries.
   auto objectLinkingLayerCreator = [&](ExecutionSession &session,
                                        const Triple &tt) {
ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
   };
 
   // Callback to inspect the cache and recompile on demand. This follows Lang's
   // LLJITWithObjectCache example.
   auto compileFunctionCreator = [&](JITTargetMachineBuilder jtmb)
       -> Expected<std::unique_ptr<IRCompileLayer::IRCompiler>> {
     if (options.jitCodeGenOptLevel)
       jtmb.setCodeGenOptLevel(*options.jitCodeGenOptLevel);
-    auto tm = jtmb.createTargetMachine();
-    if (!tm)
-      return tm.takeError();
-    return std::make_unique<TMOwningSimpleCompiler>(std::move(*tm),
+    return std::make_unique<TMOwningSimpleCompiler>(std::move(tm),
                                                     engine->cache.get());
   };
 
   // Create the LLJIT by calling the LLJITBuilder with 2 callbacks.
   auto jit =
       cantFail(llvm::orc::LLJITBuilder()
                    .setCompileFunctionCreator(compileFunctionCreator)
                    .setObjectLinkingLayerCreator(objectLinkingLayerCreator)
+                   .setDataLayout(dataLayout)
                    .create());
 
   // Add a ThreadSafemodule to the engine and return.
   ThreadSafeModule tsm(std::move(llvmModule), std::move(ctx));
   if (options.transformer)
     cantFail(tsm.withModuleDo(
         [&](llvm::Module &module) { return options.transformer(&module); }));
   cantFail(jit->addIRModule(std::move(tsm)));

mlir/lib/ExecutionEngine/JitRunner.cpp

 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/LegacyPassNameParser.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
 #include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/ToolOutputFile.h"
 #include <cstdint>
 #include <numeric>
-#include <utility>
 #include <optional>
+#include <utility>
+
+#define DEBUG_TYPE "jit-runner"
 
 using namespace mlir;
 using llvm::Error;
 
 namespace {
 /// This options struct prevents the need for global static initializers, and
 /// is only initialized if the JITRunner is invoked.
 struct Options {
llvm::cl::opt<bool> optO1{"O1",
                             llvm::cl::cat(optFlags)};
   llvm::cl::opt<bool> optO2{"O2",
                             llvm::cl::desc("Run opt passes and codegen at O2"),
                             llvm::cl::cat(optFlags)};
   llvm::cl::opt<bool> optO3{"O3",
                             llvm::cl::desc("Run opt passes and codegen at O3"),
                             llvm::cl::cat(optFlags)};
 
+  llvm::cl::list<std::string> mAttrs{
+      "mattr", llvm::cl::MiscFlags::CommaSeparated,
+      llvm::cl::desc("Target specific attributes (-mattr=help for details)"),
+      llvm::cl::value_desc("a1,+a2,-a3,..."), llvm::cl::cat(optFlags)};
+
+  llvm::cl::opt<std::string> mArch{
+      "march",
+      llvm::cl::desc("Architecture to generate code for (see --version)")};
+
   llvm::cl::OptionCategory clOptionsCategory{"linking options"};
   llvm::cl::list<std::string> clSharedLibs{
       "shared-libs", llvm::cl::desc("Libraries to link dynamically"),
       llvm::cl::MiscFlags::CommaSeparated, llvm::cl::cat(clOptionsCategory)};
 
   /// CLI variables for debugging.
   llvm::cl::opt<bool> dumpObjectFile{
       "dump-object-file",
if (flag) {
       optLevel = j;
       break;
     }
   }
   return optLevel;
 }
 
 // JIT-compile the given module and run "entryPoint" with "args" as arguments.
-static Error compileAndExecute(Options &options, Operation *module,
-                               StringRef entryPoint,
-                               CompileAndExecuteConfig config, void **args) {
+static Error
+compileAndExecute(Options &options, Operation *module, StringRef entryPoint,
+                  CompileAndExecuteConfig config, void **args,
+                  std::unique_ptr<llvm::TargetMachine> tm = nullptr) {
   std::optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel;
   if (auto clOptLevel = getCommandLineOptLevel(options))
     jitCodeGenOptLevel = static_cast<llvm::CodeGenOpt::Level>(*clOptLevel);
 
   // If shared library implements custom mlir-runner library init and destroy
   // functions, we'll use them to register the library with the execution
   // engine. Otherwise we'll pass library directly to the execution engine.
   SmallVector<SmallString<256>, 4> libPaths;
compileAndExecute(Options &options, Operation *module, StringRef entryPoint,
 
   mlir::ExecutionEngineOptions engineOptions;
   engineOptions.llvmModuleBuilder = config.llvmModuleBuilder;
   if (config.transformer)
     engineOptions.transformer = config.transformer;
   engineOptions.jitCodeGenOptLevel = jitCodeGenOptLevel;
   engineOptions.sharedLibPaths = executionEngineLibs;
   engineOptions.enableObjectDump = true;
-  auto expectedEngine = mlir::ExecutionEngine::create(module, engineOptions);
+  auto expectedEngine =
+      mlir::ExecutionEngine::create(module, engineOptions, std::move(tm));
   if (!expectedEngine)
     return expectedEngine.takeError();
 
   auto engine = std::move(*expectedEngine);
   engine->registerSymbols(runtimeSymbolMap);
 
   auto expectedFPtr = engine->lookupPacked(entryPoint);
   if (!expectedFPtr)
compileAndExecute(Options &options, Operation *module, StringRef entryPoint,
 
   // Run all dynamic library destroy callbacks to prepare for the shutdown.
   for (MlirRunnerDestroyFn destroy : destroyFns)
     destroy();
 
   return Error::success();
 }
 
-static Error compileAndExecuteVoidFunction(Options &options, Operation *module,
-                                           StringRef entryPoint,
-                                           CompileAndExecuteConfig config) {
+static Error compileAndExecuteVoidFunction(
+    Options &options, Operation *module, StringRef entryPoint,
+    CompileAndExecuteConfig config, std::unique_ptr<llvm::TargetMachine> tm) {
   auto mainFunction = dyn_cast_or_null<LLVM::LLVMFuncOp>(
       SymbolTable::lookupSymbolIn(module, entryPoint));
   if (!mainFunction || mainFunction.empty())
     return makeStringError("entry point not found");
   void *empty = nullptr;
   return compileAndExecute(options, module, entryPoint, std::move(config),
-                           &empty);
+                           &empty, std::move(tm));
 }
 
 template <typename Type>
 Error checkCompatibleReturnType(LLVM::LLVMFuncOp mainFunction);
 template <>
 Error checkCompatibleReturnType<int32_t>(LLVM::LLVMFuncOp mainFunction) {
   auto resultType = mainFunction.getFunctionType()
                         .cast<LLVM::LLVMFunctionType>()
Error checkCompatibleReturnType<float>(LLVM::LLVMFuncOp mainFunction) {
   if (!mainFunction.getFunctionType()
            .cast<LLVM::LLVMFunctionType>()
            .getReturnType()
            .isa<Float32Type>())
     return makeStringError("only single f32 function result supported");
   return Error::success();
 }
 template <typename Type>
-Error compileAndExecuteSingleReturnFunction(Options &options, Operation *module,
-                                            StringRef entryPoint,
-                                            CompileAndExecuteConfig config) {
+Error compileAndExecuteSingleReturnFunction(
+    Options &options, Operation *module, StringRef entryPoint,
+    CompileAndExecuteConfig config, std::unique_ptr<llvm::TargetMachine> tm) {
   auto mainFunction = dyn_cast_or_null<LLVM::LLVMFuncOp>(
       SymbolTable::lookupSymbolIn(module, entryPoint));
   if (!mainFunction || mainFunction.isExternal())
     return makeStringError("entry point not found");
 
   if (mainFunction.getFunctionType()
           .cast<LLVM::LLVMFunctionType>()
           .getNumParams() != 0)
     return makeStringError("function inputs not supported");
 
   if (Error error = checkCompatibleReturnType<Type>(mainFunction))
     return error;
 
   Type res;
   struct {
     void *data;
   } data;
   data.data = &res;
-  if (auto error = compileAndExecute(options, module, entryPoint,
-                                     std::move(config), (void **)&data))
+  if (auto error =
+          compileAndExecute(options, module, entryPoint, std::move(config),
+                            (void **)&data, std::move(tm)))
     return error;
 
   // Intentional printing of the output so we can test.
   llvm::outs() << res << '\n';
 
   return Error::success();
 }
 
 /// Entry point for all CPU runners. Expects the common argc/argv arguments for
 /// standard C++ main functions.
 int mlir::JitRunnerMain(int argc, char **argv, const DialectRegistry &registry,
                         JitRunnerConfig config) {
+  llvm::ExitOnError exitOnErr;
+
   // Create the options struct containing the command line options for the
   // runner. This must come before the command line options are parsed.
   Options options;
   llvm::cl::ParseCommandLineOptions(argc, argv, "MLIR CPU execution driver\n");
 
   if (options.hostSupportsJit) {
     auto j = llvm::orc::LLJITBuilder().create();
     if (j)
       llvm::outs() << "true\n";
     else {
-      llvm::consumeError(j.takeError());
       llvm::outs() << "false\n";
+      exitOnErr(j.takeError());
     }
     return 0;
   }
 
   std::optional<unsigned> optLevel = getCommandLineOptLevel(options);
   SmallVector<std::reference_wrapper<llvm::cl::opt<bool>>, 4> optFlags{
       options.optO0, options.optO1, options.optO2, options.optO3};
 
if (config.mlirTransformer)
     if (failed(config.mlirTransformer(m.get(), runnerOptions)))
       return EXIT_FAILURE;
 
   auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
   if (!tmBuilderOrError) {
     llvm::errs() << "Failed to create a JITTargetMachineBuilder for the host\n";
     return EXIT_FAILURE;
   }
+
+  // Configure TargetMachine builder based on the command line options
+  llvm::SubtargetFeatures features;
+  if (!options.mAttrs.empty()) {
+    for (unsigned i = 0; i != options.mAttrs.size(); ++i)
+      features.AddFeature(options.mAttrs[i]);
+    tmBuilderOrError->addFeatures(features.getFeatures());
+  }
+
+  if (!options.mArch.empty()) {
+    tmBuilderOrError->getTargetTriple().setArchName(options.mArch);
+  }
+
+  // Build TargetMachine
   auto tmOrError = tmBuilderOrError->createTargetMachine();
+
   if (!tmOrError) {
     llvm::errs() << "Failed to create a TargetMachine for the host\n";
-    return EXIT_FAILURE;
+    exitOnErr(tmOrError.takeError());
   }
 
+  LLVM_DEBUG({
+    llvm::dbgs() << "  JITTargetMachineBuilder is "
+                 << llvm::orc::JITTargetMachineBuilderPrinter(*tmBuilderOrError,
+                                                              "\n");
+  });
+
   CompileAndExecuteConfig compileAndExecuteConfig;
   if (optLevel) {
     compileAndExecuteConfig.transformer = mlir::makeOptimizingTransformer(
         *optLevel, /*sizeLevel=*/0, /*targetMachine=*/tmOrError->get());
   }
   compileAndExecuteConfig.llvmModuleBuilder = config.llvmModuleBuilder;
   compileAndExecuteConfig.runtimeSymbolMap = config.runtimesymbolMap;
 
   // Get the function used to compile and execute the module.
   using CompileAndExecuteFnT =
-      Error (*)(Options &, Operation *, StringRef, CompileAndExecuteConfig);
+      Error (*)(Options &, Operation *, StringRef, CompileAndExecuteConfig,
+                std::unique_ptr<llvm::TargetMachine> tm);
   auto compileAndExecuteFn =
       StringSwitch<CompileAndExecuteFnT>(options.mainFuncType.getValue())
           .Case("i32", compileAndExecuteSingleReturnFunction<int32_t>)
           .Case("i64", compileAndExecuteSingleReturnFunction<int64_t>)
           .Case("f32", compileAndExecuteSingleReturnFunction<float>)
           .Case("void", compileAndExecuteVoidFunction)
           .Default(nullptr);
 
   Error error = compileAndExecuteFn
-                    ? compileAndExecuteFn(options, m.get(),
-                                          options.mainFuncName.getValue(),
-                                          compileAndExecuteConfig)
+                    ? compileAndExecuteFn(
+                          options, m.get(), options.mainFuncName.getValue(),
+                          compileAndExecuteConfig, std::move(tmOrError.get()))
                     : makeStringError("unsupported function type");
 
   int exitCode = EXIT_SUCCESS;
   llvm::handleAllErrors(std::move(error),
                         [&exitCode](const llvm::ErrorInfoBase &info) {
                           llvm::errs() << "Error: ";
                           info.log(llvm::errs());
                           llvm::errs() << '\n';
                           exitCode = EXIT_FAILURE;
                         });
 
   return exitCode;
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/concatenate_dim_1.mlir

 // DEFINE: %{option} = enable-runtime-library=true
+// DEFINE: %{run_option} =
 // DEFINE: %{compile} = mlir-opt %s --sparse-compiler=%{option}
 // DEFINE: %{run} = mlir-cpu-runner \
 // DEFINE:  -e entry -entry-point-result=void  \
-// DEFINE:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext,%mlir_lib_dir/libmlir_runner_utils%shlibext | \
+// DEFINE:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext,%mlir_lib_dir/libmlir_runner_utils%shlibext %{run_option} | \
 // DEFINE: FileCheck %s
 //
 // RUN: %{compile} | %{run}
 //
 // Do the same run, but now with direct IR generation.
 // REDEFINE: %{option} = "enable-runtime-library=false enable-buffer-initialization=true"
 // RUN: %{compile} | %{run}
 //
-// Do the same run, but now with direct IR generation and vectorization.
-// REDEFINE: %{option} = "enable-runtime-library=false enable-buffer-initialization=true vl=2 reassociate-fp-reductions=true enable-index-optimizations=true"
+// Do the same run, but now with direct IR generation and vectorization. Enable
+// Arm SVE if supported.
+// REDEFINE: %{option} = "enable-runtime-library=false enable-buffer-initialization=true vl=4 enable-arm-sve=%ENABLE_VLA reassociate-fp-reductions=true enable-index-optimizations=true"
+// REDEFINE: %{run_option} = "%VLA_ARCH_ATTR_OPTIONS"
 // RUN: %{compile} | %{run}
 
-// Do the same run, but now with direct IR generation and, if available, VLA
-// vectorization.
-// REDEFINE: %{option} = "enable-runtime-library=false vl=4 enable-arm-sve=%ENABLE_VLA"
-// REDEFINE: %{run} = %lli \
-// REDEFINE:   --entry-function=entry_lli \
-// REDEFINE:   --extra-module=%S/Inputs/main_for_lli.ll \
-// REDEFINE:   %VLA_ARCH_ATTR_OPTIONS \
-// REDEFINE:   --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext --dlopen=%mlir_lib_dir/libmlir_runner_utils%shlibext | \
-// REDEFINE: FileCheck %s
-// RUN: %{compile} | mlir-translate -mlir-to-llvmir | %{run}
-
 #MAT_C_C = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
 #MAT_D_C = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"]}>
 #MAT_C_D = #sparse_tensor.encoding<{dimLevelType = ["compressed", "dense"]}>
 #MAT_D_D = #sparse_tensor.encoding<{
   dimLevelType = ["dense", "dense"],
   dimOrdering = affine_map<(i,j) -> (j,i)>
 }>
 

mlir/test/mlir-cpu-runner/verify-flags.mlir

+// RUN: mlir-cpu-runner %s --debug-only=jit-runner -mattr=+foo_bar -e entry -entry-point-result=void 2>&1 | FileCheck %s --check-prefixes=MATTR
+// RUN: not mlir-cpu-runner %s --debug-only=jit-runner -march=bar_foo -e entry -entry-point-result=void 2>&1 | FileCheck %s --check-prefixes=MARCH
+
+// Verify that command line args do affect the configuration
+
+// MATTR: Features = 
+// MATTR-SAME: +foo_bar
+
+// MARCH: Failed to create a TargetMachine for the host
+// MARCH-NEXT: No available targets are compatible with triple "bar_foo-{{.*}}"
+
+llvm.func @entry() -> () {
+  llvm.return
+}