[Mlir-commits] [mlir] [MLIR] split-input-file for mlir-runner (PR #184763)
Peter Žužek
llvmlistbot at llvm.org
Thu Mar 5 02:10:23 PST 2026
https://github.com/ProGTX updated https://github.com/llvm/llvm-project/pull/184763
>From ee3f3094b5227f35571cbf8290138716ade0cb98 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Peter=20=C5=BDu=C5=BEek?= <peterzuzek at gmail.com>
Date: Thu, 5 Mar 2026 11:07:19 +0100
Subject: [PATCH] [MLIR] split-input-file for mlir-runner
This patch enables running multiple separate entry points in mlir-runner
using a single input file.
Can simplify certain tests, demonstrated by updating an existing test
and adding a new one.
Current limitations:
1. All executions share the same flags (entry point name, return type, etc.)
2. No output separator.
3. No multithreading.
Closes #170255
---
mlir/lib/ExecutionEngine/JitRunner.cpp | 140 ++++++++++--------
.../Dialect/Math/CPU/mathtofuncs_ctlz.mlir | 41 ++---
.../mlir-runner/simple-split-input-file.mlir | 45 ++++++
3 files changed, 142 insertions(+), 84 deletions(-)
create mode 100644 mlir/test/mlir-runner/simple-split-input-file.mlir
diff --git a/mlir/lib/ExecutionEngine/JitRunner.cpp b/mlir/lib/ExecutionEngine/JitRunner.cpp
index db0516533afcb..05505e285c7ee 100644
--- a/mlir/lib/ExecutionEngine/JitRunner.cpp
+++ b/mlir/lib/ExecutionEngine/JitRunner.cpp
@@ -23,6 +23,7 @@
#include "mlir/IR/MLIRContext.h"
#include "mlir/Parser/Parser.h"
#include "mlir/Support/FileUtilities.h"
+#include "mlir/Support/ToolUtilities.h"
#include "mlir/Tools/ParseUtilities.h"
#include "llvm/ADT/STLExtras.h"
@@ -111,6 +112,17 @@ struct Options {
llvm::cl::desc(
"Disable implicit addition of a top-level module op during parsing"),
llvm::cl::init(false)};
+
+ llvm::cl::opt<std::string> splitInputFile{
+ "split-input-file", llvm::cl::ValueOptional,
+ llvm::cl::callback([&](const std::string &str) {
+ // Implicit value: use default marker if flag was used without value.
+ if (str.empty())
+ splitInputFile.setValue(kDefaultSplitMarker);
+ }),
+ llvm::cl::desc("Split the input file into chunks using the given or "
+ "default marker and process each chunk independently"),
+ llvm::cl::init("")};
};
struct CompileAndExecuteConfig {
@@ -131,30 +143,6 @@ struct CompileAndExecuteConfig {
} // namespace
-static OwningOpRef<Operation *> parseMLIRInput(StringRef inputFilename,
- bool insertImplicitModule,
- MLIRContext *context) {
- // Set up the input file.
- std::string errorMessage;
- auto file = openInputFile(inputFilename, &errorMessage);
- if (!file) {
- llvm::errs() << errorMessage << "\n";
- return nullptr;
- }
-
- auto sourceMgr = std::make_shared<llvm::SourceMgr>();
- sourceMgr->AddNewSourceBuffer(std::move(file), SMLoc());
- OwningOpRef<Operation *> module =
- parseSourceFileForTool(sourceMgr, context, insertImplicitModule);
- if (!module)
- return nullptr;
- if (!module.get()->hasTrait<OpTrait::SymbolTable>()) {
- llvm::errs() << "Error: top-level op must be a symbol table.\n";
- return nullptr;
- }
- return module;
-}
-
static inline Error makeStringError(const Twine &message) {
return llvm::make_error<llvm::StringError>(message.str(),
llvm::inconvertibleErrorCode());
@@ -303,6 +291,10 @@ static Error compileAndExecuteSingleReturnFunction(
return Error::success();
}
+static inline int asMainReturnCode(LogicalResult r) {
+ return r.succeeded() ? EXIT_SUCCESS : EXIT_FAILURE;
+}
+
/// 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 ®istry,
@@ -322,27 +314,9 @@ int mlir::JitRunnerMain(int argc, char **argv, const DialectRegistry ®istry,
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};
-
- MLIRContext context(registry);
-
- auto m = parseMLIRInput(options.inputFilename, !options.noImplicitModule,
- &context);
- if (!m) {
- llvm::errs() << "could not parse the input IR\n";
- return 1;
+ return EXIT_SUCCESS;
}
- JitRunnerOptions runnerOptions{options.mainFuncName, options.mainFuncType};
- 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";
@@ -376,6 +350,7 @@ int mlir::JitRunnerMain(int argc, char **argv, const DialectRegistry ®istry,
});
CompileAndExecuteConfig compileAndExecuteConfig;
+ std::optional<unsigned> optLevel = getCommandLineOptLevel(options);
if (optLevel) {
compileAndExecuteConfig.transformer = mlir::makeOptimizingTransformer(
*optLevel, /*sizeLevel=*/0, /*targetMachine=*/tmOrError->get());
@@ -394,21 +369,68 @@ int mlir::JitRunnerMain(int argc, char **argv, const DialectRegistry ®istry,
.Case("f32", compileAndExecuteSingleReturnFunction<float>)
.Case("void", compileAndExecuteVoidFunction)
.Default(nullptr);
+ if (!compileAndExecuteFn) {
+ llvm::errs() << "unsupported function type\n";
+ return EXIT_FAILURE;
+ }
+
+ std::string errorMessage;
+ auto inputFile = openInputFile(options.inputFilename, &errorMessage);
+ if (!inputFile) {
+ llvm::errs() << errorMessage << "\n";
+ return EXIT_FAILURE;
+ }
- Error error = compileAndExecuteFn
- ? 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;
+ const auto parseAndRunChunk =
+ [&](std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,
+ llvm::MemoryBufferRef sourceBuffer, [[maybe_unused]] raw_ostream &) {
+ // Tell sourceMgr about this buffer,
+ // which is what the parser will pick up.
+ auto sourceMgr = std::make_shared<llvm::SourceMgr>();
+ // Add the original buffer to the source manager to use for determining
+ // locations.
+ sourceMgr->AddNewSourceBuffer(
+ llvm::MemoryBuffer::getMemBuffer(sourceBuffer,
+ /*RequiresNullTerminator=*/false),
+ SMLoc());
+ sourceMgr->AddNewSourceBuffer(std::move(chunkBuffer), SMLoc());
+
+ MLIRContext context{registry};
+ OwningOpRef<Operation *> m = parseSourceFileForTool(
+ sourceMgr, &context, !options.noImplicitModule);
+ if (!m) {
+ llvm::errs() << "could not parse the input IR\n";
+ return failure();
+ }
+ if (!m.get()->hasTrait<OpTrait::SymbolTable>()) {
+ llvm::errs() << "Error: top-level op must be a symbol table.\n";
+ return failure();
+ }
+
+ JitRunnerOptions runnerOptions{options.mainFuncName,
+ options.mainFuncType};
+ if (config.mlirTransformer)
+ if (failed(config.mlirTransformer(m.get(), runnerOptions)))
+ return failure();
+
+ Error error = compileAndExecuteFn(
+ options, m.get(), options.mainFuncName.getValue(),
+ compileAndExecuteConfig, std::move(tmOrError.get()));
+
+ auto chunkSuccess = success();
+ llvm::handleAllErrors(std::move(error),
+ [&chunkSuccess](const llvm::ErrorInfoBase &info) {
+ llvm::errs() << "Error: ";
+ info.log(llvm::errs());
+ llvm::errs() << '\n';
+ chunkSuccess = failure();
+ });
+ return chunkSuccess;
+ };
+
+ return asMainReturnCode(splitAndProcessBuffer(
+ llvm::MemoryBuffer::getMemBuffer(inputFile->getMemBufferRef(),
+ /*RequiresNullTerminator=*/false),
+ parseAndRunChunk, llvm::nulls(), options.splitInputFile,
+ options.splitInputFile));
}
diff --git a/mlir/test/Integration/Dialect/Math/CPU/mathtofuncs_ctlz.mlir b/mlir/test/Integration/Dialect/Math/CPU/mathtofuncs_ctlz.mlir
index cd9483410138f..7e7661ceb3648 100644
--- a/mlir/test/Integration/Dialect/Math/CPU/mathtofuncs_ctlz.mlir
+++ b/mlir/test/Integration/Dialect/Math/CPU/mathtofuncs_ctlz.mlir
@@ -1,48 +1,39 @@
// RUN: mlir-opt %s \
+// RUN: --split-input-file \
// RUN: -pass-pipeline="builtin.module( \
// RUN: convert-math-to-funcs{convert-ctlz}, \
// RUN: func.func(convert-scf-to-cf,convert-arith-to-llvm), \
// RUN: convert-func-to-llvm, \
// RUN: convert-cf-to-llvm, \
// RUN: reconcile-unrealized-casts)" \
-// RUN: | mlir-runner -e test_7i32_to_29 -entry-point-result=i32 | FileCheck %s --check-prefix=CHECK_TEST_7i32_TO_29
+// RUN: | mlir-runner --split-input-file -entry-point-result=i32 \
+// RUN: | FileCheck %s
-func.func @test_7i32_to_29() -> i32 {
+func.func @main() -> i32 {
%arg = arith.constant 7 : i32
%0 = math.ctlz %arg : i32
func.return %0 : i32
}
-// CHECK_TEST_7i32_TO_29: 29
+// CHECK: 29
-// RUN: mlir-opt %s \
-// RUN: -pass-pipeline="builtin.module( \
-// RUN: convert-math-to-funcs{convert-ctlz}, \
-// RUN: func.func(convert-scf-to-cf,convert-arith-to-llvm), \
-// RUN: convert-func-to-llvm, \
-// RUN: convert-cf-to-llvm, \
-// RUN: reconcile-unrealized-casts)" \
-// RUN: | mlir-runner -e test_zero -entry-point-result=i32 | FileCheck %s --check-prefix=CHECK_TEST_ZERO
+// -----
-func.func @test_zero() -> i32 {
+func.func @main() -> i32 {
%arg = arith.constant 0 : i32
%0 = math.ctlz %arg : i32
func.return %0 : i32
}
-// CHECK_TEST_ZERO: 32
+// CHECK: 32
-// Apparently mlir-runner doesn't support i8 return values, so testing i64 instead
-// RUN: mlir-opt %s \
-// RUN: -pass-pipeline="builtin.module( \
-// RUN: convert-math-to-funcs, \
-// RUN: func.func(convert-scf-to-cf,convert-arith-to-llvm), \
-// RUN: convert-func-to-llvm, \
-// RUN: convert-cf-to-llvm, \
-// RUN: reconcile-unrealized-casts)" \
-// RUN: | mlir-runner -e test_7i64_to_61 -entry-point-result=i64 | FileCheck %s --check-prefix=CHECK_TEST_7i64_TO_61
+// -----
-func.func @test_7i64_to_61() -> i64 {
+
+// Apparently mlir-runner doesn't support i8 return values, so testing i64 instead
+func.func @main() -> i32 {
%arg = arith.constant 7 : i64
%0 = math.ctlz %arg : i64
- func.return %0 : i64
+ // We can safely truncate 64-bit result to 32 bits
+ %1 = arith.trunci %0 : i64 to i32
+ func.return %1 : i32
}
-// CHECK_TEST_7i64_TO_61: 61
+// CHECK: 61
diff --git a/mlir/test/mlir-runner/simple-split-input-file.mlir b/mlir/test/mlir-runner/simple-split-input-file.mlir
new file mode 100644
index 0000000000000..047f42583fdec
--- /dev/null
+++ b/mlir/test/mlir-runner/simple-split-input-file.mlir
@@ -0,0 +1,45 @@
+// RUN: mlir-runner %s --split-input-file %if target={{s390x-.*}} %{ -argext-abi-check=false %} \
+// RUN: | FileCheck %s
+
+// Declarations of C library functions.
+llvm.func @logbf(f32) -> f32
+
+// Check that a simple function with a nested call works.
+llvm.func @main() -> f32 {
+ %0 = llvm.mlir.constant(-4.200000e+02 : f32) : f32
+ %1 = llvm.call @logbf(%0) : (f32) -> f32
+ llvm.return %1 : f32
+}
+// CHECK: 8.000000e+00
+
+// -----
+
+// Declarations of C library functions.
+llvm.func @malloc(i64) -> !llvm.ptr
+llvm.func @free(!llvm.ptr)
+
+// Helper typed functions wrapping calls to "malloc" and "free".
+llvm.func @allocation() -> !llvm.ptr {
+ %0 = llvm.mlir.constant(4 : index) : i64
+ %1 = llvm.call @malloc(%0) : (i64) -> !llvm.ptr
+ llvm.return %1 : !llvm.ptr
+}
+llvm.func @deallocation(%arg0: !llvm.ptr) {
+ llvm.call @free(%arg0) : (!llvm.ptr) -> ()
+ llvm.return
+}
+
+// Check that allocation and deallocation works, and that a custom entry point
+// works.
+llvm.func @main() -> f32 {
+ %0 = llvm.call @allocation() : () -> !llvm.ptr
+ %1 = llvm.mlir.constant(0 : index) : i64
+ %2 = llvm.mlir.constant(1.234000e+03 : f32) : f32
+ %3 = llvm.getelementptr %0[%1] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+ llvm.store %2, %3 : f32, !llvm.ptr
+ %4 = llvm.getelementptr %0[%1] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+ %5 = llvm.load %4 : !llvm.ptr -> f32
+ llvm.call @deallocation(%0) : (!llvm.ptr) -> ()
+ llvm.return %5 : f32
+}
+// CHECK: 1.234000e+03
More information about the Mlir-commits
mailing list