Adding configurable compiler plugins for Bazel builds.

MODULE.bazel

@@ -67,7 +67,8 @@ use_repo(
     "vulkan_sdk",
 )
 
-# IREE extension creates llvm-raw (when IREE is root) and other IREE-specific repos
+# IREE extension creates MLIR-adjacent source repos when IREE is root and other
+# IREE-specific repos unconditionally.
 iree_ext = use_extension("//build_tools/bazel:extensions.bzl", "iree_extension")
 use_repo(
     iree_ext,
@@ -83,6 +84,7 @@ use_repo(
     "rccl",
     "spirv_cross",
     "stablehlo",
+    "torch-mlir-raw",
     "tracy_client",
     "vulkan_headers",
     "webgpu_headers",
@@ -103,3 +105,12 @@ llvm_configure(
         "X86",
     ],
 )
+
+# Configure torch-mlir using its upstream Bazel overlay. The configured
+# repository is only needed when the Torch input plugin is selected.
+torch_mlir_configure = use_repo_rule("@torch-mlir-raw//utils/bazel:configure.bzl", "torch_mlir_configure")
+
+torch_mlir_configure(
+    name = "torch-mlir",
+    src_workspace = "@torch-mlir-raw//:CMakeLists.txt",
+)

build_tools/bazel/BZLMOD_LLVM.md

@@ -1,7 +1,8 @@
-# Using IREE with a Custom LLVM via Bzlmod
+# Using IREE with Custom MLIR-Adjacent Dependencies via Bzlmod
 
-This document explains how projects that depend on IREE can provide their own LLVM
-instead of using IREE's bundled version.
+This document explains how projects that depend on IREE can provide their own
+LLVM, StableHLO, torch-mlir, and compiler plugin registry instead of using
+IREE's bundled defaults.
 
 ## Terminology
 
@@ -38,9 +39,11 @@ use_repo(ext, "repo_a", "repo_b")
 Imports a repository rule from another module so it can be called directly in
 MODULE.bazel to create a repository.
 
-### `llvm-raw`
-A repository containing the raw LLVM source code. This is the input to the LLVM
-build configuration.
+### Raw source repositories
+Repositories such as `llvm-raw` and `torch-mlir-raw` contain unconfigured
+upstream source trees. They are inputs to repository rules that overlay Bazel
+BUILD files and produce configured repositories such as `llvm-project` and
+`torch-mlir`.
 
 ### `llvm-project`
 The configured LLVM repository created by `llvm_configure`. It overlays Bazel BUILD
@@ -53,8 +56,8 @@ The bzlmod module name for LLVM's Bazel integration (located at
 
 ## How It Works
 
-IREE's module extension (`iree_extension`) creates the `llvm-raw` repository
-**only when IREE is the root module**:
+IREE's module extension (`iree_extension`) creates MLIR-adjacent source
+repositories **only when IREE is the root module**:
 
 ```python
 # In build_tools/bazel/extensions.bzl
@@ -65,11 +68,21 @@ def _iree_extension_impl(module_ctx):
             build_file_content = "# empty",
             path = "third_party/llvm-project",
         )
+        local_repository(
+            name = "stablehlo",
+            path = "third_party/stablehlo",
+        )
+        new_local_repository(
+            name = "torch-mlir-raw",
+            build_file_content = "# empty - BUILD files overlaid by torch_mlir_configure",
+            path = "third_party/torch-mlir",
+        )
     # ... other repos
 ```
 
 When your project depends on IREE, IREE is **not** the root module - your project is.
-Therefore, IREE's extension will not create `llvm-raw`, and you must provide it yourself.
+Therefore, IREE's extension will not create these MLIR-adjacent repositories,
+and you must provide the ones needed by the compiler plugins you enable.
 
 ## MODULE.bazel Ordering
 
@@ -82,7 +95,7 @@ The order of statements in MODULE.bazel matters:
 5. `use_repo()` - must come after its corresponding `use_extension()`
 6. `use_repo_rule()` + invocation - can reference repos created by earlier extensions
 
-## Example: Using Your Own LLVM
+## Example: Using Your Own LLVM, StableHLO, and torch-mlir
 
 ```python
 # my_project/MODULE.bazel
@@ -114,7 +127,7 @@ local_path_override(
     path = "my/custom/llvm-project/utils/bazel",
 )
 
-# Create your own llvm-raw repository pointing to your LLVM
+# Create your own raw repositories pointing to the upstream projects you want.
 new_local_repository = use_repo_rule(
     "@bazel_tools//tools/build_defs/repo:local.bzl",
     "new_local_repository",
@@ -124,6 +137,19 @@ new_local_repository(
     path = "my/custom/llvm-project",
     build_file_content = "# empty",
 )
+new_local_repository(
+    name = "torch-mlir-raw",
+    path = "my/custom/torch-mlir",
+    build_file_content = "# empty",
+)
+local_repository = use_repo_rule(
+    "@bazel_tools//tools/build_defs/repo:local.bzl",
+    "local_repository",
+)
+local_repository(
+    name = "stablehlo",
+    path = "my/custom/stablehlo",
+)
 
 # Use LLVM's extension for third-party deps (gmp, mpfr, etc.)
 llvm_repos_ext = use_extension(
@@ -161,8 +187,34 @@ llvm_configure = use_repo_rule(
     "llvm_configure",
 )
 llvm_configure(name = "llvm-project")
+
+# Configure torch-mlir from your raw source repository if you enable the Torch
+# input plugin.
+torch_mlir_configure = use_repo_rule(
+    "@torch-mlir-raw//utils/bazel:configure.bzl",
+    "torch_mlir_configure",
+)
+torch_mlir_configure(
+    name = "torch-mlir",
+    src_workspace = "@torch-mlir-raw//:CMakeLists.txt",
+)
+```
+
+## Custom Compiler Plugin Registry
+
+IREE's `//compiler/plugins` package loads the plugin registry from the root
+workspace:
+
+```python
+load("@//build_tools/bazel:default_compiler_plugins.bzl", ...)
 ```
 
+When IREE is the root module, this resolves to IREE's default registry. A
+downstream root workspace can provide a file at the same path to register
+additional compiler plugins, replace registration targets, or change the
+default enabled plugin IDs. In-tree IREE plugin labels should be qualified with
+`@iree_core//...` from such a downstream file.
+
 ## Using LLVM from an HTTP Archive
 
 If you want to fetch LLVM from a release tarball instead of a local path:
@@ -204,7 +256,15 @@ When providing your own LLVM, ensure compatibility with IREE:
 
 ### "repository 'llvm-raw' is not defined"
 You haven't created the `llvm-raw` repository. As the root module, you must
-define it yourself (see examples above).
+define it yourself if you configure LLVM (see examples above).
+
+### "repository 'stablehlo' is not defined"
+You enabled the StableHLO input plugin without providing a `stablehlo`
+repository from your root module.
+
+### "repository 'torch-mlir' is not defined"
+You enabled the Torch input plugin without configuring a `torch-mlir` repository
+from your root module.
 
 ### Build errors in LLVM code
 Your LLVM version may be incompatible with IREE. Check that your LLVM commit

build_tools/bazel/default_compiler_plugins.bzl

@@ -0,0 +1,49 @@
+# Copyright 2026 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+"""Default compiler plugin registry for Bazel builds.
+
+IREE's plugin aggregation BUILD file loads this file from the root workspace
+with @//build_tools/bazel:default_compiler_plugins.bzl. When IREE is built as
+part of a larger Bazel workspace, that root workspace may provide a file at the
+same path to register its own compiler plugins or override the default set.
+"""
+
+def iree_default_compiler_plugin_ids():
+    """Returns plugin IDs enabled by plain Bazel invocation."""
+    return [
+        "input_stablehlo",
+        "input_tosa",
+        "hal_target_cuda",
+        "hal_target_llvm_cpu",
+        "hal_target_local",
+        "hal_target_metal_spirv",
+        "hal_target_rocm",
+        "hal_target_vmvx",
+        "hal_target_vulkan_spirv",
+        "example",
+        "simple_io_sample",
+    ]
+
+def iree_default_compiler_plugins():
+    """Returns all known compiler plugin registration targets."""
+    return {
+        # Input plugins.
+        "input_stablehlo": "//compiler/plugins/input/StableHLO:registration",
+        "input_tosa": "//compiler/plugins/input/TOSA:registration",
+        "input_torch": "//compiler/plugins/input/Torch:registration",
+        # Target plugins.
+        "hal_target_cuda": "//compiler/plugins/target/CUDA",
+        "hal_target_llvm_cpu": "//compiler/plugins/target/LLVMCPU",
+        "hal_target_local": "//compiler/plugins/target/Local",
+        "hal_target_metal_spirv": "//compiler/plugins/target/MetalSPIRV",
+        "hal_target_rocm": "//compiler/plugins/target/ROCM",
+        "hal_target_vmvx": "//compiler/plugins/target/VMVX",
+        "hal_target_vulkan_spirv": "//compiler/plugins/target/VulkanSPIRV",
+        # Sample plugins.
+        "example": "//samples/compiler_plugins/example:registration",
+        "simple_io_sample": "//samples/compiler_plugins/simple_io_sample:registration",
+    }

build_tools/bazel/extensions.bzl

@@ -13,14 +13,26 @@ load("//build_tools/bazel:workspace.bzl", "cuda_auto_configure")
 def _iree_extension_impl(module_ctx):
     """Implementation of the IREE module extension."""
 
-    # Create llvm-raw only when IREE is the root module.
-    # This allows downstream consumers to provide their own LLVM.
-    if any([m.is_root and m.name == "iree_core" for m in module_ctx.modules]):
+    iree_is_root = any([m.is_root and m.name == "iree_core" for m in module_ctx.modules])
+
+    # Create MLIR-adjacent source repositories only when IREE is the root
+    # module. Downstream consumers can provide their own LLVM, StableHLO, and
+    # torch-mlir repositories when IREE is part of a larger project.
+    if iree_is_root:
         new_local_repository(
             name = "llvm-raw",
             build_file_content = "# empty",
             path = "third_party/llvm-project",
         )
+        local_repository(
+            name = "stablehlo",
+            path = "third_party/stablehlo",
+        )
+        new_local_repository(
+            name = "torch-mlir-raw",
+            build_file_content = "# empty - BUILD files overlaid by torch_mlir_configure",
+            path = "third_party/torch-mlir",
+        )
 
     # Googletest
     local_repository(
@@ -42,12 +54,6 @@ def _iree_extension_impl(module_ctx):
         path = "third_party/vulkan_headers",
     )
 
-    # StableHLO
-    local_repository(
-        name = "stablehlo",
-        path = "third_party/stablehlo",
-    )
-
     # Benchmark
     local_repository(
         name = "com_google_benchmark",

build_tools/bazel/iree.bazelrc

@@ -48,6 +48,8 @@ build --force_pic
 # Build settings flag aliases. See https://bazel.build/rules/config
 ###############################################################################
 build --flag_alias=iree_drivers=//runtime/src/iree/hal/drivers:enabled_drivers
+build --flag_alias=iree_runtime_drivers=//runtime/src/iree/hal/drivers:enabled_drivers
+build --flag_alias=iree_compiler_plugins=//compiler/plugins:enabled_plugins
 build --flag_alias=iree_link_compiler_shared=//compiler/src/iree/compiler/API:link_shared
 # Equivalent to CMake flag IREE_ENABLE_RUNTIME_TRACING + IREE_TRACING_PROVIDER.
 # Builds the runtime with tracing support enabled.

build_tools/bazel/workspace.bzl

@@ -8,6 +8,7 @@
 load("@bazel_skylib//lib:paths.bzl", "paths")
 load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe")
 
+IREE_INPUT_TORCH_ENV_KEY = "IREE_INPUT_TORCH"
 CUDA_TOOLKIT_ROOT_ENV_KEY = "IREE_CUDA_TOOLKIT_ROOT"
 
 # Our CI docker images use a stripped down CUDA directory tree in some
@@ -67,6 +68,78 @@ cuda_auto_configure = repository_rule(
     },
 )
 
+def torch_mlir_auto_configure_impl(repository_ctx):
+    """Conditionally configures torch-mlir based on IREE_INPUT_TORCH env var."""
+    env = repository_ctx.os.environ
+    iree_repo_alias = repository_ctx.attr.iree_repo_alias
+    enabled = env.get(IREE_INPUT_TORCH_ENV_KEY, "OFF").upper() in ["ON", "TRUE", "1", "YES"]
+
+    if enabled:
+        # Run torch-mlir's configure to create the overlay.
+        # We need to find the torch-mlir source and run its overlay script.
+        torch_mlir_path = repository_ctx.path(
+            Label("%s//:third_party/torch-mlir/CMakeLists.txt" % iree_repo_alias),
+        ).dirname
+        bazel_path = torch_mlir_path.get_child("utils").get_child("bazel")
+        overlay_path = bazel_path.get_child("torch-mlir-overlay")
+        script_path = bazel_path.get_child("overlay_directories.py")
+
+        python_bin = repository_ctx.which("python3")
+        if not python_bin:
+            python_bin = repository_ctx.which("python")
+        if not python_bin:
+            fail("Failed to find python3 binary for torch-mlir configuration")
+
+        cmd = [
+            python_bin,
+            script_path,
+            "--src",
+            torch_mlir_path,
+            "--overlay",
+            overlay_path,
+            "--target",
+            ".",
+        ]
+        exec_result = repository_ctx.execute(cmd, timeout = 60)
+
+        if exec_result.return_code != 0:
+            fail(("Failed to configure torch-mlir: '{cmd}'\n" +
+                  "Exited with code {return_code}\n" +
+                  "stdout:\n{stdout}\n" +
+                  "stderr:\n{stderr}\n").format(
+                cmd = " ".join([str(arg) for arg in cmd]),
+                return_code = exec_result.return_code,
+                stdout = exec_result.stdout,
+                stderr = exec_result.stderr,
+            ))
+    else:
+        # Create stub repository when torch-mlir is disabled.
+        repository_ctx.file(
+            "BUILD.bazel",
+            content = """# Stub: torch-mlir disabled (IREE_INPUT_TORCH != ON)
+package(default_visibility = ["//visibility:public"])
+
+# Provide empty targets that dependent code can reference.
+# These will fail at build time if actually used.
+""",
+        )
+
+torch_mlir_auto_configure = repository_rule(
+    environ = [IREE_INPUT_TORCH_ENV_KEY],
+    implementation = torch_mlir_auto_configure_impl,
+    local = True,
+    attrs = {
+        "iree_repo_alias": attr.string(default = "@iree_core"),
+    },
+)
+
+def configure_iree_torch_mlir_deps(iree_repo_alias = None):
+    maybe(
+        torch_mlir_auto_configure,
+        name = "torch-mlir",
+        iree_repo_alias = iree_repo_alias,
+    )
+
 def configure_iree_cuda_deps(iree_repo_alias = None):
     maybe(
         cuda_auto_configure,