Add support for Bazel 6 & 7

build_tools/utils.bzl

@@ -0,0 +1,26 @@
+# Copyright 2024 The StableHLO Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A collection of starlark functions for use in BUILD files that are useful for StableHLO."""
+
+def is_bzlmod_enabled():
+    """Determine whether bzlmod mode is enabled."""
+
+    # If bzlmod is enabled, then `str(Label(...))` returns a canonical label,
+    # these start with `@@`.
+    return str(Label("//:invalid")).startswith("@@")
+
+def workspace_name():
+    """Return the name of the workspace."""
+    return "_main" if is_bzlmod_enabled() else "stablehlo"

stablehlo/conversions/linalg/tests/BUILD.bazel

@@ -25,9 +25,15 @@ expand_template(
     testonly = True,
     out = "lit.site.cfg.py",
     substitutions = {
-        "\"@STABLEHLO_SOURCE_DIR@\"": "os.path.join(os.environ['TEST_SRCDIR'], '_main')",
-        "\"@STABLEHLO_TOOLS_DIR@\"": "os.path.join(os.environ['TEST_SRCDIR'], '_main')",
-        "@LIT_SITE_CFG_IN_HEADER@": "# Autogenerated, do not edit.",
+        "\"@STABLEHLO_SOURCE_DIR@\"": "RUNFILES_DIR",
+        "\"@STABLEHLO_TOOLS_DIR@\"": "RUNFILES_DIR",
+        "@LIT_SITE_CFG_IN_HEADER@": '''# Autogenerated, do not edit.
+from python.runfiles import Runfiles
+from pathlib import Path
+
+r = Runfiles.Create()
+LITE_CFG_PY = Path(r.Rlocation("stablehlo/stablehlo/conversions/linalg/tests/lit.cfg.py"))
+RUNFILES_DIR = LITE_CFG_PY.parents[4].absolute().as_posix()''',
         "@LLVM_TOOLS_DIR@": "../llvm-project/llvm",
     },
     template = "lit.site.cfg.py.in",
@@ -46,6 +52,7 @@ expand_template(
             "@llvm-project//llvm:FileCheck",
         ],
         tags = ["stablehlo_linalg_tests"],
+        deps = ["@rules_python//python/runfiles"],
     )
     for src in glob(["**/*.mlir"])
 ]

stablehlo/conversions/tosa/tests/BUILD.bazel

@@ -25,9 +25,15 @@ expand_template(
     testonly = True,
     out = "lit.site.cfg.py",
     substitutions = {
-        "\"@STABLEHLO_SOURCE_DIR@\"": "os.path.join(os.environ['TEST_SRCDIR'], '_main')",
-        "\"@STABLEHLO_TOOLS_DIR@\"": "os.path.join(os.environ['TEST_SRCDIR'], '_main')",
-        "@LIT_SITE_CFG_IN_HEADER@": "# Autogenerated, do not edit.",
+        "\"@STABLEHLO_SOURCE_DIR@\"": "RUNFILES_DIR",
+        "\"@STABLEHLO_TOOLS_DIR@\"": "RUNFILES_DIR",
+        "@LIT_SITE_CFG_IN_HEADER@": '''# Autogenerated, do not edit.
+from python.runfiles import Runfiles
+from pathlib import Path
+
+r = Runfiles.Create()
+LITE_CFG_PY = Path(r.Rlocation("stablehlo/stablehlo/conversions/tosa/tests/lit.cfg.py"))
+RUNFILES_DIR = LITE_CFG_PY.parents[4].absolute().as_posix()''',
         "@LLVM_TOOLS_DIR@": "../llvm-project/llvm",
     },
     template = "lit.site.cfg.py.in",
@@ -46,6 +52,7 @@ expand_template(
             "@llvm-project//llvm:FileCheck",
         ],
         tags = ["stablehlo_tosa_tests"],
+        deps = ["@rules_python//python/runfiles"],
     )
     for src in glob(["**/*.mlir"])
 ]

stablehlo/testdata/BUILD.bazel

@@ -25,9 +25,15 @@ expand_template(
     testonly = True,
     out = "lit.site.cfg.py",
     substitutions = {
-        "\"@STABLEHLO_SOURCE_DIR@\"": "os.path.join(os.environ['TEST_SRCDIR'], '_main')",
-        "\"@STABLEHLO_TOOLS_DIR@\"": "os.path.join(os.environ['TEST_SRCDIR'], '_main')",
-        "@LIT_SITE_CFG_IN_HEADER@": "# Autogenerated, do not edit.",
+        "\"@STABLEHLO_SOURCE_DIR@\"": "RUNFILES_DIR",
+        "\"@STABLEHLO_TOOLS_DIR@\"": "RUNFILES_DIR",
+        "@LIT_SITE_CFG_IN_HEADER@": '''# Autogenerated, do not edit.
+from python.runfiles import Runfiles
+from pathlib import Path
+
+r = Runfiles.Create()
+LITE_CFG_PY = Path(r.Rlocation("stablehlo/stablehlo/testdata/lit.cfg.py"))
+RUNFILES_DIR = LITE_CFG_PY.parents[2].absolute().as_posix()''',
         "@LLVM_TOOLS_DIR@": "../llvm-project/llvm",
     },
     template = "lit.site.cfg.py.in",
@@ -47,6 +53,7 @@ expand_template(
             "@llvm-project//llvm:FileCheck",
         ],
         tags = ["stablehlo_testdata_tests"],
+        deps = ["@rules_python//python/runfiles"],
     )
     for src in glob(["**/*.mlir"])
 ]

stablehlo/tests/BUILD.bazel

@@ -134,9 +134,15 @@ expand_template(
     testonly = True,
     out = "lit.site.cfg.py",
     substitutions = {
-        "\"@STABLEHLO_SOURCE_DIR@\"": "os.path.join(os.environ['TEST_SRCDIR'], '_main')",
-        "\"@STABLEHLO_TOOLS_DIR@\"": "os.path.join(os.environ['TEST_SRCDIR'], '_main')",
-        "@LIT_SITE_CFG_IN_HEADER@": "# Autogenerated, do not edit.",
+        "\"@STABLEHLO_SOURCE_DIR@\"": "RUNFILES_DIR",
+        "\"@STABLEHLO_TOOLS_DIR@\"": "RUNFILES_DIR",
+        "@LIT_SITE_CFG_IN_HEADER@": '''# Autogenerated, do not edit.
+from python.runfiles import Runfiles
+from pathlib import Path
+
+r = Runfiles.Create()
+LITE_CFG_PY = Path(r.Rlocation("stablehlo/stablehlo/tests/lit.cfg.py"))
+RUNFILES_DIR = LITE_CFG_PY.parents[2].absolute().as_posix()''',
         "@LLVM_TOOLS_DIR@": "../llvm-project/llvm",
     },
     template = "lit.site.cfg.py.in",
@@ -157,6 +163,7 @@ expand_template(
             "@llvm-project//llvm:not",
         ] + glob(["%s.bc" % src]),
         tags = ["stablehlo_tests"],
+        deps = ["@rules_python//python/runfiles"],
     )
     for src in glob(["**/*.mlir"])
 ]