Skip to content

Commit 34e24d7

Browse files
committed
Add tuning spec flags to TuningSpecOptions for C API support
Signed-off-by: Bangtian Liu <[email protected]>
1 parent 4cc4576 commit 34e24d7

File tree

13 files changed

+205
-14
lines changed

13 files changed

+205
-14
lines changed

compiler/src/iree/compiler/API/Internal/BUILD.bazel

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@ iree_compiler_cc_library(
2323
],
2424
deps = [
2525
"//compiler/bindings/c:headers",
26+
"//compiler/src/iree/compiler/Codegen/Common:PassHeaders",
27+
"//compiler/src/iree/compiler/Codegen/Common:TuningSpecOptions",
2628
"//compiler/src/iree/compiler/ConstEval",
2729
"//compiler/src/iree/compiler/Dialect/VM/Target:init_targets",
2830
"//compiler/src/iree/compiler/Dialect/VM/Target/Bytecode",

compiler/src/iree/compiler/API/Internal/CMakeLists.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ iree_cc_library(
3030
MLIRParser
3131
MLIRRemarkStreamer
3232
MLIRSupport
33+
iree::compiler::Codegen::Common::PassHeaders
34+
iree::compiler::Codegen::Common::TuningSpecOptions
3335
iree::compiler::ConstEval
3436
iree::compiler::Dialect::VM::Target::Bytecode
3537
iree::compiler::Dialect::VM::Target::C

compiler/src/iree/compiler/API/Internal/CompilerDriver.cpp

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@
4141
#include <limits>
4242

4343
#include "iree/compiler/API/Internal/Diagnostics.h"
44+
#include "iree/compiler/Codegen/Common/Options.h"
45+
#include "iree/compiler/Codegen/Common/Passes.h"
4446
#include "iree/compiler/ConstEval/Passes.h"
4547
#include "iree/compiler/Dialect/VM/Target/init_targets.h"
4648
#include "iree/compiler/Dialect/VM/Transforms/Passes.h"
@@ -250,6 +252,7 @@ struct GlobalInit {
250252
ParameterOptions *clParameterOptions = nullptr;
251253
DispatchCreationOptions *clDispatchCreationOptions = nullptr;
252254
SchedulingOptions *clSchedulingOptions = nullptr;
255+
TuningSpecOptions *clTuningSpecOptions = nullptr;
253256
IREE::HAL::TargetOptions *clHalTargetOptions = nullptr;
254257
IREE::VM::TargetOptions *clVmTargetOptions = nullptr;
255258
IREE::VM::BytecodeTargetOptions *clBytecodeTargetOptions = nullptr;
@@ -297,6 +300,7 @@ void GlobalInit::registerCommandLineOptions() {
297300
clParameterOptions = &ParameterOptions::FromFlags::get();
298301
clDispatchCreationOptions = &DispatchCreationOptions::FromFlags::get();
299302
clSchedulingOptions = &SchedulingOptions::FromFlags::get();
303+
clTuningSpecOptions = &TuningSpecOptions::FromFlags::get();
300304
clHalTargetOptions = &IREE::HAL::TargetOptions::FromFlags::get();
301305
clVmTargetOptions = &IREE::VM::TargetOptions::FromFlags::get();
302306
clBytecodeTargetOptions = &IREE::VM::BytecodeTargetOptions::FromFlags::get();
@@ -409,6 +413,7 @@ struct Session {
409413
GlobalOptimizationOptions highLevelOptimizationOptions;
410414
DispatchCreationOptions dispatchCreationOptions;
411415
SchedulingOptions schedulingOptions;
416+
TuningSpecOptions tuningSpecOptions;
412417
IREE::HAL::TargetOptions halTargetOptions;
413418
IREE::VM::TargetOptions vmTargetOptions;
414419
IREE::VM::BytecodeTargetOptions bytecodeTargetOptions;
@@ -438,6 +443,7 @@ Session::Session(GlobalInit &globalInit)
438443
parameterOptions = *globalInit.clParameterOptions;
439444
dispatchCreationOptions = *globalInit.clDispatchCreationOptions;
440445
schedulingOptions = *globalInit.clSchedulingOptions;
446+
tuningSpecOptions = *globalInit.clTuningSpecOptions;
441447
halTargetOptions = *globalInit.clHalTargetOptions;
442448
vmTargetOptions = *globalInit.clVmTargetOptions;
443449
bytecodeTargetOptions = *globalInit.clBytecodeTargetOptions;
@@ -460,6 +466,7 @@ Session::Session(GlobalInit &globalInit)
460466
parameterOptions.bindOptions(binder);
461467
dispatchCreationOptions.bindOptions(binder);
462468
schedulingOptions.bindOptions(binder);
469+
tuningSpecOptions.bindOptions(binder);
463470
halTargetOptions.bindOptions(binder);
464471
vmTargetOptions.bindOptions(binder);
465472
bytecodeTargetOptions.bindOptions(binder);
@@ -1000,6 +1007,18 @@ bool Invocation::runPipeline(enum iree_compiler_pipeline_t pipeline) {
10001007
}
10011008
});
10021009

1010+
// Only pipelines that use MaterializeTuningSpecsPass need tuning spec
1011+
// options. STD and HAL_EXECUTABLE both run buildHALTransformPassPipeline
1012+
// which includes ConfigureExecutablesPass that adds
1013+
// MaterializeTuningSpecsPass.
1014+
bool needsTuningSpecs = (pipeline == IREE_COMPILER_PIPELINE_STD ||
1015+
pipeline == IREE_COMPILER_PIPELINE_HAL_EXECUTABLE);
1016+
if (needsTuningSpecs) {
1017+
setGlobalTuningSpecOptions(&session.tuningSpecOptions);
1018+
}
1019+
auto resetTuningSpecOptions =
1020+
llvm::scope_exit([&]() { setGlobalTuningSpecOptions(nullptr); });
1021+
10031022
switch (pipeline) {
10041023
case IREE_COMPILER_PIPELINE_STD: {
10051024
IREEVMPipelinePhase compileFrom;
@@ -1089,10 +1108,13 @@ bool Invocation::runTextualPassPipeline(const char *textPassPipeline) {
10891108
llvm::errs()))) {
10901109
return false;
10911110
}
1092-
if (failed(passManager->run(parsedModule))) {
1093-
return false;
1094-
}
1095-
return true;
1111+
1112+
// Inject tuning spec options before pipeline execution.
1113+
setGlobalTuningSpecOptions(&session.tuningSpecOptions);
1114+
auto resetTuningSpecOptions =
1115+
llvm::scope_exit([&]() { setGlobalTuningSpecOptions(nullptr); });
1116+
1117+
return succeeded(passManager->run(parsedModule));
10961118
}
10971119

10981120
Error *Invocation::outputIR(Output &output) {

compiler/src/iree/compiler/API/test/BUILD.bazel

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,3 +21,13 @@ iree_compiler_cc_test(
2121
"@llvm-project//mlir:CAPIIRHeaders",
2222
],
2323
)
24+
25+
iree_compiler_cc_test(
26+
name = "tuning-spec-flags-test",
27+
testonly = True,
28+
srcs = ["tuning_spec_flags_test.c"],
29+
deps = [
30+
"//compiler/bindings/c:headers",
31+
"//compiler/src/iree/compiler/API:Impl",
32+
],
33+
)

compiler/src/iree/compiler/API/test/CMakeLists.txt

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,16 @@ iree_cc_test(
2121
iree::compiler::bindings::c::headers
2222
)
2323

24+
iree_cc_test(
25+
NAME
26+
tuning-spec-flags-test
27+
SRCS
28+
"tuning_spec_flags_test.c"
29+
DEPS
30+
iree::compiler::API::Impl
31+
iree::compiler::bindings::c::headers
32+
)
33+
2434
### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
2535

2636
# Move to bin/ directory and systematically name more appropriately.
Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
// Copyright 2026 The IREE Authors
2+
//
3+
// Licensed under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
7+
// Test that verifies the tuning spec path flag can be set via C API.
8+
// Before the TuningSpecOptions refactoring, this flag would be rejected
9+
// because it was not registered in the Session's local OptionsBinder.
10+
// Other tuning spec flags (enable-default-tuning-specs, dump-tuning-specs-to)
11+
// remain CLI-only and are not tested here.
12+
13+
#include <stdio.h>
14+
#include "iree/compiler/embedding_api.h"
15+
16+
int main(int argc, char **argv) {
17+
ireeCompilerGlobalInitialize();
18+
iree_compiler_session_t *session = ireeCompilerSessionCreate();
19+
20+
const char *flags[] = {"--iree-codegen-tuning-spec-path=/tmp/spec.mlir"};
21+
22+
iree_compiler_error_t *err = ireeCompilerSessionSetFlags(session, 1, flags);
23+
if (err) {
24+
fprintf(stderr, "FAIL: Tuning spec path flag not accepted: %s\n",
25+
ireeCompilerErrorGetMessage(err));
26+
ireeCompilerErrorDestroy(err);
27+
ireeCompilerSessionDestroy(session);
28+
ireeCompilerGlobalShutdown();
29+
return 1;
30+
}
31+
32+
printf("PASS: Tuning spec path flag accepted via C API\n");
33+
34+
ireeCompilerSessionDestroy(session);
35+
ireeCompilerGlobalShutdown();
36+
return 0;
37+
}

compiler/src/iree/compiler/Codegen/Common/BUILD.bazel

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,17 @@ iree_compiler_cc_library(
5050
],
5151
)
5252

53+
iree_compiler_cc_library(
54+
name = "TuningSpecOptions",
55+
srcs = ["Options.cpp"],
56+
hdrs = ["Options.h"],
57+
deps = [
58+
"//compiler/src/iree/compiler/Utils",
59+
"@llvm-project//llvm:Support",
60+
"@llvm-project//mlir:IR",
61+
],
62+
)
63+
5364
iree_gentbl_cc_library(
5465
name = "FoldTensorExtractOpIncGen",
5566
tbl_outs = [
@@ -194,6 +205,7 @@ iree_compiler_cc_library(
194205
deps = [
195206
":PassHeaders",
196207
":PassesIncGen",
208+
":TuningSpecOptions",
197209
"//compiler/src/iree/compiler/Codegen/Common:FoldTensorExtractOpIncGen",
198210
"//compiler/src/iree/compiler/Codegen/Dialect/CPU/IR:IREECPUDialect",
199211
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",

compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,20 @@ iree_cc_library(
3939
PUBLIC
4040
)
4141

42+
iree_cc_library(
43+
NAME
44+
TuningSpecOptions
45+
HDRS
46+
"Options.h"
47+
SRCS
48+
"Options.cpp"
49+
DEPS
50+
LLVMSupport
51+
MLIRIR
52+
iree::compiler::Utils
53+
PUBLIC
54+
)
55+
4256
iree_tablegen_library(
4357
NAME
4458
FoldTensorExtractOpIncGen
@@ -171,6 +185,7 @@ iree_cc_library(
171185
DEPS
172186
::PassHeaders
173187
::PassesIncGen
188+
::TuningSpecOptions
174189
IREELinalgTransformDialect
175190
LLVMSupport
176191
MLIRAMDGPUDialect

compiler/src/iree/compiler/Codegen/Common/MaterializeTuningSpecsPass.cpp

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
66

77
#include <cassert>
8+
#include "iree/compiler/Codegen/Common/Options.h"
89
#include "iree/compiler/Codegen/Common/Passes.h"
910
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
1011
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
@@ -42,12 +43,6 @@ namespace mlir::iree_compiler {
4243

4344
namespace {
4445

45-
llvm::cl::opt<std::string> clCodegenTuningSpecPath(
46-
"iree-codegen-tuning-spec-path",
47-
llvm::cl::desc("File path to a module containing a tuning spec (transform "
48-
"dialect library)."),
49-
llvm::cl::init(""));
50-
5146
llvm::cl::opt<bool> clCodegenEnableDefaultTuningSpecs(
5247
"iree-codegen-enable-default-tuning-specs",
5348
llvm::cl::desc("Whether to enable default tuning spec transform libraries "
@@ -63,6 +58,19 @@ llvm::cl::opt<std::string> clCodegenTuningSpecDumpDir(
6358

6459
using mlir::transform::NamedSequenceOp;
6560

61+
// Global accessor for tuning spec options.
62+
// This is set by the compiler driver before pipeline execution and cleared
63+
// after completion.
64+
static TuningSpecOptions *globalTuningSpecOptions = nullptr;
65+
66+
static const TuningSpecOptions &getTuningSpecOptions() {
67+
if (globalTuningSpecOptions) {
68+
return *globalTuningSpecOptions;
69+
}
70+
// Fallback to FromFlags for backward compatibility with standalone usage.
71+
return TuningSpecOptions::FromFlags::get();
72+
}
73+
6674
static LogicalResult dumpFinalTuningSpecToDir(ModuleOp tuningSpec) {
6775
StringRef dir = clCodegenTuningSpecDumpDir;
6876
if (dir.empty()) {
@@ -97,16 +105,17 @@ static LogicalResult dumpFinalTuningSpecToDir(ModuleOp tuningSpec) {
97105

98106
static FailureOr<ModuleOp>
99107
getUserTuningSpec(ModuleOp module, IREE::Codegen::IREECodegenDialect &dialect) {
100-
if (clCodegenTuningSpecPath.empty()) {
108+
const std::string &tuningSpecPath = getTuningSpecOptions().tuningSpecPath;
109+
if (tuningSpecPath.empty()) {
101110
return failure();
102111
}
103112

104113
FailureOr<ModuleOp> maybeTransformLibrary =
105-
dialect.getOrLoadTransformLibraryModule(clCodegenTuningSpecPath);
114+
dialect.getOrLoadTransformLibraryModule(tuningSpecPath);
106115
if (failed(maybeTransformLibrary)) {
107116
return module->emitError()
108117
<< "Failed to load tuning spec transform dialect library from "
109-
<< clCodegenTuningSpecPath;
118+
<< tuningSpecPath;
110119
}
111120

112121
return *maybeTransformLibrary;
@@ -187,7 +196,7 @@ struct MaterializeTuningSpecsPass final
187196

188197
FailureOr<ModuleOp> userTuningSpec = getUserTuningSpec(moduleOp, *dialect);
189198
const bool hasUserTuningSpec = succeeded(userTuningSpec);
190-
if (!hasUserTuningSpec && !clCodegenTuningSpecPath.empty()) {
199+
if (!hasUserTuningSpec && !getTuningSpecOptions().tuningSpecPath.empty()) {
191200
// When a user spec is requested but fails to load, this is a hard
192201
// failure.
193202
return signalPassFailure();
@@ -279,4 +288,9 @@ struct MaterializeTuningSpecsPass final
279288
};
280289

281290
} // namespace
291+
292+
void setGlobalTuningSpecOptions(TuningSpecOptions *options) {
293+
globalTuningSpecOptions = options;
294+
}
295+
282296
} // namespace mlir::iree_compiler
Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
// Copyright 2026 The IREE Authors
2+
//
3+
// Licensed under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
7+
#include "iree/compiler/Codegen/Common/Options.h"
8+
9+
IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::TuningSpecOptions);
10+
11+
namespace mlir::iree_compiler {
12+
13+
void TuningSpecOptions::bindOptions(OptionsBinder &binder) {
14+
static llvm::cl::OptionCategory category(
15+
"IREE codegen tuning spec options",
16+
"Options for controlling codegen tuning spec loading and configuration.");
17+
18+
binder.opt<std::string>(
19+
"iree-codegen-tuning-spec-path", tuningSpecPath,
20+
llvm::cl::desc("File path to a module containing a tuning spec "
21+
"(transform dialect library)."),
22+
llvm::cl::cat(category));
23+
}
24+
25+
} // namespace mlir::iree_compiler

0 commit comments

Comments
 (0)