Skip to content

Commit c562954

Browse files
benvanikbenvanik
committed
[Bazel] Configure static compiler plugins
Replace the hardcoded Bazel StaticLinkedPlugins.inc contents with a generated registry driven by //compiler/plugins:enabled_plugins and the --iree_compiler_plugins alias. This mirrors the runtime driver selection model while preserving the existing default plugin set for plain Bazel invocation. Teach iree-bazel-configure/configure_bazel.py to resolve IREE_COMPILER_PLUGINS, including all/exclusion syntax, from buildable local inputs. The wrapper defaults to all buildable plugins so developer builds pick up optional inputs such as Torch when the submodule is initialized, while explicit plugin lists fail on unknown or unavailable entries. CUDA selection also records the repository env var consumed by @iree_cuda so plugin selection and external repository configuration agree. Add Bazel wiring for the Torch input plugin and conditional @torch-mlir repository overlay. This lets Bazel-built tools link input_torch and process Torch dialect input through the same conversion paths covered by the existing lit tests. The Torch InputConversion CMake files are now generated from BUILD.bazel through explicit torch-mlir target mappings; the root Torch CMake file remains hand-maintained because it owns the CMake-only torch-mlir vendor setup and plugin registration. MetalSPIRV stays in the buildable default set because it is a compiler target, not Darwin-only runtime support.
1 parent e380f43 commit c562954

15 files changed

Lines changed: 774 additions & 59 deletions

File tree

‎MODULE.bazel‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ use_repo(
8383
"rccl",
8484
"spirv_cross",
8585
"stablehlo",
86+
"torch-mlir",
8687
"tracy_client",
8788
"vulkan_headers",
8889
"webgpu_headers",

‎build_tools/bazel/extensions.bzl‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88

99
load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
1010
load("@bazel_tools//tools/build_defs/repo:local.bzl", "local_repository", "new_local_repository")
11-
load("//build_tools/bazel:workspace.bzl", "cuda_auto_configure")
11+
load("//build_tools/bazel:workspace.bzl", "cuda_auto_configure", "torch_mlir_auto_configure")
1212

1313
def _iree_extension_impl(module_ctx):
1414
"""Implementation of the IREE module extension."""
@@ -118,6 +118,12 @@ def _iree_extension_impl(module_ctx):
118118
iree_repo_alias = "@iree_core",
119119
)
120120

121+
# torch-mlir auto-configuration (conditional on IREE_INPUT_TORCH)
122+
torch_mlir_auto_configure(
123+
name = "torch-mlir",
124+
iree_repo_alias = "@iree_core",
125+
)
126+
121127
iree_extension = module_extension(
122128
implementation = _iree_extension_impl,
123129
)

‎build_tools/bazel/iree.bazelrc‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@ build --force_pic
4848
# Build settings flag aliases. See https://bazel.build/rules/config
4949
###############################################################################
5050
build --flag_alias=iree_drivers=//runtime/src/iree/hal/drivers:enabled_drivers
51+
build --flag_alias=iree_runtime_drivers=//runtime/src/iree/hal/drivers:enabled_drivers
52+
build --flag_alias=iree_compiler_plugins=//compiler/plugins:enabled_plugins
5153
build --flag_alias=iree_link_compiler_shared=//compiler/src/iree/compiler/API:link_shared
5254
# Equivalent to CMake flag IREE_ENABLE_RUNTIME_TRACING + IREE_TRACING_PROVIDER.
5355
# Builds the runtime with tracing support enabled.

‎build_tools/bazel/workspace.bzl‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
load("@bazel_skylib//lib:paths.bzl", "paths")
99
load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe")
1010

11+
IREE_INPUT_TORCH_ENV_KEY = "IREE_INPUT_TORCH"
1112
CUDA_TOOLKIT_ROOT_ENV_KEY = "IREE_CUDA_TOOLKIT_ROOT"
1213

1314
# Our CI docker images use a stripped down CUDA directory tree in some
@@ -67,6 +68,78 @@ cuda_auto_configure = repository_rule(
6768
},
6869
)
6970

71+
def torch_mlir_auto_configure_impl(repository_ctx):
72+
"""Conditionally configures torch-mlir based on IREE_INPUT_TORCH env var."""
73+
env = repository_ctx.os.environ
74+
iree_repo_alias = repository_ctx.attr.iree_repo_alias
75+
enabled = env.get(IREE_INPUT_TORCH_ENV_KEY, "OFF").upper() in ["ON", "TRUE", "1", "YES"]
76+
77+
if enabled:
78+
# Run torch-mlir's configure to create the overlay.
79+
# We need to find the torch-mlir source and run its overlay script.
80+
torch_mlir_path = repository_ctx.path(
81+
Label("%s//:third_party/torch-mlir/CMakeLists.txt" % iree_repo_alias),
82+
).dirname
83+
bazel_path = torch_mlir_path.get_child("utils").get_child("bazel")
84+
overlay_path = bazel_path.get_child("torch-mlir-overlay")
85+
script_path = bazel_path.get_child("overlay_directories.py")
86+
87+
python_bin = repository_ctx.which("python3")
88+
if not python_bin:
89+
python_bin = repository_ctx.which("python")
90+
if not python_bin:
91+
fail("Failed to find python3 binary for torch-mlir configuration")
92+
93+
cmd = [
94+
python_bin,
95+
script_path,
96+
"--src",
97+
torch_mlir_path,
98+
"--overlay",
99+
overlay_path,
100+
"--target",
101+
".",
102+
]
103+
exec_result = repository_ctx.execute(cmd, timeout = 60)
104+
105+
if exec_result.return_code != 0:
106+
fail(("Failed to configure torch-mlir: '{cmd}'\n" +
107+
"Exited with code {return_code}\n" +
108+
"stdout:\n{stdout}\n" +
109+
"stderr:\n{stderr}\n").format(
110+
cmd = " ".join([str(arg) for arg in cmd]),
111+
return_code = exec_result.return_code,
112+
stdout = exec_result.stdout,
113+
stderr = exec_result.stderr,
114+
))
115+
else:
116+
# Create stub repository when torch-mlir is disabled.
117+
repository_ctx.file(
118+
"BUILD.bazel",
119+
content = """# Stub: torch-mlir disabled (IREE_INPUT_TORCH != ON)
120+
package(default_visibility = ["//visibility:public"])
121+
122+
# Provide empty targets that dependent code can reference.
123+
# These will fail at build time if actually used.
124+
""",
125+
)
126+
127+
torch_mlir_auto_configure = repository_rule(
128+
environ = [IREE_INPUT_TORCH_ENV_KEY],
129+
implementation = torch_mlir_auto_configure_impl,
130+
local = True,
131+
attrs = {
132+
"iree_repo_alias": attr.string(default = "@iree_core"),
133+
},
134+
)
135+
136+
def configure_iree_torch_mlir_deps(iree_repo_alias = None):
137+
maybe(
138+
torch_mlir_auto_configure,
139+
name = "torch-mlir",
140+
iree_repo_alias = iree_repo_alias,
141+
)
142+
70143
def configure_iree_cuda_deps(iree_repo_alias = None):
71144
maybe(
72145
cuda_auto_configure,

‎build_tools/bazel_to_cmake/bazel_to_cmake_targets.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ def __init__(self, repo_map: Dict[str, str]):
1414
self._repo_map = repo_map
1515

1616
iree_core_repo = self._repo_alias("@iree_core")
17+
torch_mlir_cmake_prefix = "iree::compiler::plugins::input::Torch"
1718
self._update_target_mappings(
1819
{
1920
# Internal utilities to emulate various binary/library options.
@@ -108,6 +109,49 @@ def __init__(self, repo_map: Dict[str, str]):
108109
"@stablehlo//:vhlo_ops": [
109110
"VhloOps",
110111
],
112+
# Torch-MLIR.
113+
#
114+
# The CMake vendor shim groups several fine-grained Bazel
115+
# overlay libraries into coarser CMake libraries.
116+
"@torch-mlir//:TorchMLIRConversionPasses": [
117+
f"{torch_mlir_cmake_prefix}::torch-mlir::ConversionPasses",
118+
],
119+
"@torch-mlir//:TorchMLIRTMTensorDialect": [
120+
f"{torch_mlir_cmake_prefix}::torch-mlir-dialects::TMTensorDialectIR",
121+
],
122+
"@torch-mlir//:TorchMLIRTorchConversionDialect": [
123+
f"{torch_mlir_cmake_prefix}::torch-mlir::TorchConversionDialectIR",
124+
],
125+
"@torch-mlir//:TorchMLIRTorchConversionPasses": [
126+
f"{torch_mlir_cmake_prefix}::torch-mlir::ConversionPasses",
127+
],
128+
"@torch-mlir//:TorchMLIRTorchConversionToMLProgram": [
129+
f"{torch_mlir_cmake_prefix}::torch-mlir::ConversionPasses",
130+
],
131+
"@torch-mlir//:TorchMLIRTorchDialect": [
132+
f"{torch_mlir_cmake_prefix}::torch-mlir::TorchDialectIR",
133+
],
134+
"@torch-mlir//:TorchMLIRTorchOnnxToTorch": [
135+
f"{torch_mlir_cmake_prefix}::torch-mlir::TorchOnnxToTorchPasses",
136+
],
137+
"@torch-mlir//:TorchMLIRTorchPasses": [
138+
f"{torch_mlir_cmake_prefix}::torch-mlir::TorchDialectPasses",
139+
],
140+
"@torch-mlir//:TorchMLIRTorchToArith": [
141+
f"{torch_mlir_cmake_prefix}::torch-mlir::ConversionPasses",
142+
],
143+
"@torch-mlir//:TorchMLIRTorchToLinalg": [
144+
f"{torch_mlir_cmake_prefix}::torch-mlir::ConversionPasses",
145+
],
146+
"@torch-mlir//:TorchMLIRTorchToSCF": [
147+
f"{torch_mlir_cmake_prefix}::torch-mlir::ConversionPasses",
148+
],
149+
"@torch-mlir//:TorchMLIRTorchToTMTensor": [
150+
f"{torch_mlir_cmake_prefix}::torch-mlir::ConversionPasses",
151+
],
152+
"@torch-mlir//:TorchMLIRTorchToTensor": [
153+
f"{torch_mlir_cmake_prefix}::torch-mlir::ConversionPasses",
154+
],
111155
# HIP
112156
"@hip_api_headers//:headers": [
113157
"hip_api_headers::headers",

‎build_tools/bin/iree-bazel-configure‎

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -48,18 +48,35 @@ ENVIRONMENT VARIABLES
4848
Values: ON/YES/TRUE/Y/1 (enabled), OFF/NO/FALSE/N/0 (disabled)
4949
Defaults: local-sync, local-task, vulkan enabled (matches CMake)
5050
51-
Note: Compiler backends are always built in Bazel (plugin architecture).
51+
Compiler Plugins:
52+
IREE_COMPILER_PLUGINS=all - All buildable plugins (default)
53+
IREE_COMPILER_PLUGINS=all,-X - All except plugin X
54+
IREE_COMPILER_PLUGINS=X,Y - Only plugins X and Y
55+
56+
Plugin detection (for "all"):
57+
- CUDA: requires IREE_CUDA_TOOLKIT_ROOT/IREE_CUDA_DEPS_DIR,
58+
CUDA_ROOT/CUDA_HOME/CUDA_PATH/CUDA_TOOLKIT_ROOT_DIR, or nvcc in PATH
59+
- ROCm: requires ROCM_PATH/ROCM_ROOT/ROCM_HOME/HIP_PATH or hipcc in PATH
60+
- MetalSPIRV: always buildable
61+
- StableHLO: requires third_party/stablehlo submodule
62+
- Torch: requires third_party/torch-mlir submodule
5263
5364
EXAMPLES
54-
# Configure with defaults
65+
# Configure with defaults (all buildable plugins)
5566
iree-bazel-configure
5667
57-
# Enable CUDA
68+
# Enable CUDA runtime driver
5869
IREE_HAL_DRIVER_CUDA=ON iree-bazel-configure
5970
60-
# Enable multiple drivers
71+
# Enable multiple runtime drivers
6172
IREE_HAL_DRIVER_CUDA=ON IREE_HAL_DRIVER_VULKAN=ON iree-bazel-configure
6273
74+
# Exclude torch plugin (faster builds)
75+
IREE_COMPILER_PLUGINS="all,-input_torch" iree-bazel-configure
76+
77+
# Only specific plugins
78+
IREE_COMPILER_PLUGINS="hal_target_llvm_cpu,hal_target_vulkan_spirv" iree-bazel-configure
79+
6380
AFTER CONFIGURATION
6481
iree-bazel-build Build IREE targets
6582
iree-bazel-test Run tests with configured drivers
@@ -84,6 +101,11 @@ iree_info "Configuring IREE for Bazel in ${IREE_BAZEL_WORKTREE_DIR}..."
84101
# Run platform detection.
85102
if [[ -f "configure_bazel.py" ]]; then
86103
iree_info "Running configure_bazel.py..."
104+
# Default to all buildable plugins unless explicitly set.
105+
# Users who want selective plugins can set IREE_COMPILER_PLUGINS explicitly.
106+
if [[ -z "${IREE_COMPILER_PLUGINS:-}" ]]; then
107+
export IREE_COMPILER_PLUGINS="all"
108+
fi
87109
python3 configure_bazel.py
88110
else
89111
iree_error "configure_bazel.py not found in ${IREE_BAZEL_WORKTREE_DIR}"

‎compiler/plugins/BUILD.bazel‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
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+
# bazel-to-cmake: skip
8+
# This file contains Bazel-specific plugin configuration with no CMake equivalent.
9+
# CMake plugin registration uses iree_compiler_register_plugin() in each plugin's
10+
# CMakeLists.txt instead.
11+
12+
load(":plugins.bzl", "iree_compiler_plugins")
13+
14+
package(default_visibility = ["//visibility:public"])
15+
16+
# All compiler plugins with their registration targets.
17+
# The macro generates:
18+
# - config_setting: {plugin_id}_enabled
19+
# - genrule: {plugin_id}_plugin_id -> {plugin_id}.inc
20+
# - filegroup: plugins_incs (all enabled .inc files)
21+
# - cc_library: plugins_deps (all enabled registration deps)
22+
iree_compiler_plugins(
23+
name = "plugins",
24+
# Plugins enabled by the plain Bazel default. configure_bazel.py may
25+
# override this list and opt into torch when torch-mlir is available.
26+
default_plugins = [
27+
"input_stablehlo",
28+
"input_tosa",
29+
"hal_target_cuda",
30+
"hal_target_llvm_cpu",
31+
"hal_target_local",
32+
"hal_target_metal_spirv",
33+
"hal_target_rocm",
34+
"hal_target_vmvx",
35+
"hal_target_vulkan_spirv",
36+
"example",
37+
"simple_io_sample",
38+
],
39+
plugins = {
40+
# Input plugins
41+
"input_stablehlo": "//compiler/plugins/input/StableHLO:registration",
42+
"input_tosa": "//compiler/plugins/input/TOSA:registration",
43+
"input_torch": "//compiler/plugins/input/Torch:registration",
44+
# Target plugins
45+
"hal_target_cuda": "//compiler/plugins/target/CUDA",
46+
"hal_target_llvm_cpu": "//compiler/plugins/target/LLVMCPU",
47+
"hal_target_local": "//compiler/plugins/target/Local",
48+
"hal_target_metal_spirv": "//compiler/plugins/target/MetalSPIRV",
49+
"hal_target_rocm": "//compiler/plugins/target/ROCM",
50+
"hal_target_vmvx": "//compiler/plugins/target/VMVX",
51+
"hal_target_vulkan_spirv": "//compiler/plugins/target/VulkanSPIRV",
52+
# Sample plugins
53+
"example": "//samples/compiler_plugins/example:registration",
54+
"simple_io_sample": "//samples/compiler_plugins/simple_io_sample:registration",
55+
},
56+
)

‎compiler/plugins/input/Torch/BUILD.bazel‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
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+
# Torch-MLIR input plugin for IREE.
8+
# This plugin is conditionally built when IREE_INPUT_TORCH=ON.
9+
#
10+
# The torch-mlir libraries are provided by @torch-mlir which is configured
11+
# by the torch_mlir_auto_configure repository rule.
12+
13+
load("//build_tools/bazel:build_defs.oss.bzl", "iree_compiler_cc_library")
14+
15+
package(
16+
default_visibility = ["//visibility:public"],
17+
features = ["layering_check"],
18+
licenses = ["notice"],
19+
)
20+
21+
iree_compiler_cc_library(
22+
name = "registration",
23+
srcs = ["PluginRegistration.cpp"],
24+
deps = [
25+
"//compiler/plugins/input/Torch/InputConversion",
26+
"//compiler/src/iree/compiler/Dialect/Flow/IR",
27+
"//compiler/src/iree/compiler/Dialect/HAL/IR",
28+
"//compiler/src/iree/compiler/Dialect/LinalgExt/IR",
29+
"//compiler/src/iree/compiler/Dialect/Stream/IR",
30+
"//compiler/src/iree/compiler/Dialect/TensorExt/IR",
31+
"//compiler/src/iree/compiler/Dialect/Util/IR",
32+
"//compiler/src/iree/compiler/PluginAPI",
33+
"@llvm-project//mlir:IR",
34+
"@llvm-project//mlir:MLProgramDialect",
35+
"@llvm-project//mlir:Pass",
36+
"@torch-mlir//:TorchMLIRConversionPasses",
37+
"@torch-mlir//:TorchMLIRTMTensorDialect",
38+
"@torch-mlir//:TorchMLIRTorchConversionDialect",
39+
"@torch-mlir//:TorchMLIRTorchConversionPasses",
40+
"@torch-mlir//:TorchMLIRTorchDialect",
41+
"@torch-mlir//:TorchMLIRTorchOnnxToTorch",
42+
"@torch-mlir//:TorchMLIRTorchPasses",
43+
],
44+
)

0 commit comments

Comments
 (0)