|
| 1 | +# Copyright 2025-2026 Ant Group Co., Ltd. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""Single-point gateway for CuTeDSL's private MLIR/NVVM bindings. |
| 16 | +
|
| 17 | +cuLA kernels are written against CuTeDSL's code-generation API. Most of that API |
| 18 | +is public (``cutlass.cutlass_dsl``, ``cutlass.cute``, ...); a small part is not: |
| 19 | +the generated MLIR dialect bindings under ``cutlass._mlir``. CuTeDSL ships them |
| 20 | +as implementation detail and provides no stability contract across patch |
| 21 | +releases -- the ``tcgen05_ld/st`` breakage between CutDSL 4.5.2 and 4.5.3 was one |
| 22 | +such incident. |
| 23 | +
|
| 24 | +This module is the ONLY place in cuLA that may import from ``cutlass._mlir``. |
| 25 | +Kernel modules bind their dialect aliases from here:: |
| 26 | +
|
| 27 | + from cula.ops._mlir_compat import arith as _arith |
| 28 | + from cula.ops._mlir_compat import ir |
| 29 | + from cula.ops._mlir_compat import llvm as _llvm |
| 30 | + from cula.ops._mlir_compat import vector as _vector |
| 31 | +
|
| 32 | +Design goals: |
| 33 | +
|
| 34 | +- Lazy: nothing is imported until a kernel actually needs a binding, so plain |
| 35 | + imports of cuLA never touch ``cutlass._mlir``. |
| 36 | +- Explicity: an out-of-contract CuTeDSL version, or a missing/renamed binding, |
| 37 | + fails fast at first use with an actionable ``RuntimeError`` instead of |
| 38 | + surfacing mid-JIT as a confusing compile error (or silently emitting a |
| 39 | + different kernel). |
| 40 | +- Zero dependencies: version parsing is done with a small regex so this module |
| 41 | + stays usable in any environment cuLA can be installed into. |
| 42 | +
|
| 43 | +The version contract mirrors ``pyproject.toml`` (including its ``!=4.5.0`` |
| 44 | +exclusion) and the canary probes make a broken binding fail fast with an |
| 45 | +actionable message instead of surfacing mid-JIT. |
| 46 | +""" |
| 47 | + |
| 48 | +from __future__ import annotations |
| 49 | + |
| 50 | +import importlib |
| 51 | +import re |
| 52 | +from typing import Any, Final |
| 53 | + |
| 54 | +# Version contract for nvidia-cutlass-dsl. Kept in sync with |
| 55 | +# ``pyproject.toml``; when CuTeDSL is bumped, extend ``_SUPPORTED_MIN`` / |
| 56 | +# ``_SUPPORTED_MAX`` only after the new release has been validated against the |
| 57 | +# canaries below (and ideally against the SM90/SM100 kernel test suites). |
| 58 | +_SUPPORTED_MIN: Final[tuple[int, ...]] = (4, 4, 2) |
| 59 | +_SUPPORTED_MAX: Final[tuple[int, ...]] = (4, 7, 0) |
| 60 | +_EXCLUDED_VERSIONS: Final[frozenset[tuple[int, ...]]] = frozenset({(4, 5, 0)}) |
| 61 | + |
| 62 | +# dialect name -> (package, attribute) inside ``cutlass``. |
| 63 | +_PRIVATE_TABLE: Final[dict[str, tuple[str, str]]] = { |
| 64 | + "arith": ("cutlass._mlir", "dialects.arith"), |
| 65 | + "cute": ("cutlass._mlir", "dialects.cute"), |
| 66 | + "ir": ("cutlass._mlir", "ir"), |
| 67 | + "llvm": ("cutlass._mlir", "dialects.llvm"), |
| 68 | + "nvvm": ("cutlass._mlir", "dialects.nvvm"), |
| 69 | + "vector": ("cutlass._mlir", "dialects.vector"), |
| 70 | +} |
| 71 | + |
| 72 | +# Canary entry points: attribute paths that must be present on each dialect |
| 73 | +# binding for cuLA's kernel code to be emitted correctly. These are the names |
| 74 | +# the migrated consumers actually call; extend the list when new usages land. |
| 75 | +_CANARIES: Final[dict[str, tuple[tuple[str, ...], ...]]] = { |
| 76 | + "arith": (("constant",),), |
| 77 | + "cute": (), |
| 78 | + "ir": (("Type", "parse"), ("VectorType", "get")), |
| 79 | + "llvm": (("inline_asm",), ("extractvalue",)), |
| 80 | + "nvvm": (), |
| 81 | + "vector": (("bitcast",), ("extract_strided_slice",)), |
| 82 | +} |
| 83 | + |
| 84 | +# Entry points whose name changed across CutDSL versions: for each group (a |
| 85 | +# tuple of variant paths), at least one variant must exist. For example |
| 86 | +# ``vector.extractelement`` was replaced by ``vector.extract`` in the 4.6 |
| 87 | +# line; cuLA helpers dispatch on whichever one exists. |
| 88 | +_ANY_OF_CANARIES: Final[dict[str, tuple[tuple[tuple[str, ...], ...], ...]]] = { |
| 89 | + "vector": ((("extract",), ("extractelement",)),), |
| 90 | +} |
| 91 | + |
| 92 | +_VERSION_RE: Final[re.Pattern[str]] = re.compile(r"^(\d+)(?:\.(\d+))?(?:\.(\d+))?(?:[a-z0-9._+-]*)?$") |
| 93 | + |
| 94 | +_INCIDENT_NOTE: Final[str] = ( |
| 95 | + "cuLA depends on CuTeDSL's private MLIR bindings (`cutlass._mlir`), which are " |
| 96 | + "implementation detail and were broken by a CuTeDSL patch release once " |
| 97 | + "already (`tcgen05_ld/st`, 4.5.2 -> 4.5.3). To avoid silent kernel " |
| 98 | + "miscompiles, cuLA refuses bindings outside its validated contract." |
| 99 | +) |
| 100 | + |
| 101 | +_CACHE: Final[dict[str, Any]] = {} |
| 102 | + |
| 103 | + |
| 104 | +def _parse_version(version: str) -> tuple[int, int, int]: |
| 105 | + """Parse ``X.Y.Z[.devN...]`` into a comparable ``(major, minor, patch)`` tuple. |
| 106 | +
|
| 107 | + Missing components normalize to zero so that ``4.5`` and ``4.7`` compare |
| 108 | + against the contract exactly like ``4.5.0`` and ``4.7.0`` (a two-component |
| 109 | + ``4.7`` must not slip under the ``< 4.7.0`` upper bound). |
| 110 | + """ |
| 111 | + match = _VERSION_RE.match(version.strip()) |
| 112 | + if match is None: |
| 113 | + raise RuntimeError( |
| 114 | + f"cuLA cannot parse the installed CuTeDSL version {version!r}; " |
| 115 | + f"refusing to use its private MLIR bindings. {_INCIDENT_NOTE}" |
| 116 | + ) |
| 117 | + major, minor, patch = (int(part) if part is not None else 0 for part in match.groups()) |
| 118 | + return (major, minor, patch) |
| 119 | + |
| 120 | + |
| 121 | +def _installed_version() -> tuple[int, int, int]: |
| 122 | + try: |
| 123 | + import cutlass # noqa: PLC0415 |
| 124 | + except ImportError: |
| 125 | + raise RuntimeError( |
| 126 | + "cuLA requires the nvidia-cutlass-dsl package; install it with `pip install 'nvidia-cutlass-dsl>=4.4.2,<4.7'`." |
| 127 | + ) from None |
| 128 | + version = getattr(cutlass, "__version__", None) |
| 129 | + if not isinstance(version, str): |
| 130 | + raise RuntimeError( |
| 131 | + f"CuTeDSL is installed but exposes no `__version__` (got {version!r}); " |
| 132 | + f"refusing to use its private MLIR bindings. {_INCIDENT_NOTE}" |
| 133 | + ) |
| 134 | + return _parse_version(version) |
| 135 | + |
| 136 | + |
| 137 | +def _check_contract(version: tuple[int, ...]) -> None: |
| 138 | + if version in _EXCLUDED_VERSIONS: |
| 139 | + raise RuntimeError( |
| 140 | + f"Installed CuTeDSL version {'.'.join(map(str, version))} is explicitly " |
| 141 | + f"excluded by cuLA (see pyproject.toml). {_INCIDENT_NOTE}" |
| 142 | + ) |
| 143 | + if not (_SUPPORTED_MIN <= version < _SUPPORTED_MAX): |
| 144 | + raise RuntimeError( |
| 145 | + f"Installed CuTeDSL version {'.'.join(map(str, version))} is outside the " |
| 146 | + f"range validated by cuLA ({'.'.join(map(str, _SUPPORTED_MIN))} to " |
| 147 | + f"{'.'.join(map(str, _SUPPORTED_MAX))}, exclusive). " |
| 148 | + f"{_INCIDENT_NOTE} To proceed, pin the validated range in " |
| 149 | + f"pyproject.toml and re-validate the canaries in this module." |
| 150 | + ) |
| 151 | + |
| 152 | + |
| 153 | +def _has_attribute_path(module: Any, path: tuple[str, ...]) -> bool: |
| 154 | + owner = module |
| 155 | + for part in path: |
| 156 | + owner = getattr(owner, part, None) |
| 157 | + if owner is None: |
| 158 | + return False |
| 159 | + return True |
| 160 | + |
| 161 | + |
| 162 | +def _load(dialect: str) -> Any: |
| 163 | + if dialect in _CACHE: |
| 164 | + return _CACHE[dialect] |
| 165 | + |
| 166 | + package, attribute = _PRIVATE_TABLE[dialect] |
| 167 | + try: |
| 168 | + module = importlib.import_module(package) |
| 169 | + except ImportError as exc: |
| 170 | + raise RuntimeError( |
| 171 | + f"Unable to import CuTeDSL's private {dialect!r} bindings ({package}): {exc}. {_INCIDENT_NOTE}" |
| 172 | + ) from exc |
| 173 | + if attribute: |
| 174 | + for part in attribute.split("."): |
| 175 | + module = getattr(module, part, None) |
| 176 | + if module is None: |
| 177 | + break |
| 178 | + if module is None: |
| 179 | + raise RuntimeError(f"CuTeDSL no longer exposes {dialect!r} bindings ({package}.{attribute}). {_INCIDENT_NOTE}") |
| 180 | + |
| 181 | + for canary in _CANARIES[dialect]: |
| 182 | + owner: Any = module |
| 183 | + for part in canary: |
| 184 | + owner = getattr(owner, part, None) |
| 185 | + if owner is None: |
| 186 | + raise RuntimeError( |
| 187 | + f"CuTeDSL dialect {dialect!r} is missing the canary entry point " |
| 188 | + f"{'.'.join(canary)} used by cuLA kernels. {_INCIDENT_NOTE}" |
| 189 | + ) |
| 190 | + for group in _ANY_OF_CANARIES.get(dialect, ()): |
| 191 | + if not any(_has_attribute_path(module, variant) for variant in group): |
| 192 | + raise RuntimeError( |
| 193 | + f"CuTeDSL dialect {dialect!r} exposes none of the entry points " |
| 194 | + f"{' / '.join('.'.join(variant) for variant in group)} expected by " |
| 195 | + f"cuLA kernels. {_INCIDENT_NOTE}" |
| 196 | + ) |
| 197 | + |
| 198 | + _CACHE[dialect] = module |
| 199 | + return module |
| 200 | + |
| 201 | + |
| 202 | +def __getattr__(name: str) -> Any: |
| 203 | + """Lazy, contract-checked access to private dialect bindings.""" |
| 204 | + if name not in _PRIVATE_TABLE: |
| 205 | + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") |
| 206 | + version = _installed_version() |
| 207 | + _check_contract(version) |
| 208 | + return _load(name) |
| 209 | + |
| 210 | + |
| 211 | +def cutlass_dsl_version() -> str | None: |
| 212 | + """Installed CuTeDSL version string, or None when not installed.""" |
| 213 | + try: |
| 214 | + import cutlass # noqa: PLC0415 |
| 215 | + except ImportError: |
| 216 | + return None |
| 217 | + version = getattr(cutlass, "__version__", None) |
| 218 | + return version if isinstance(version, str) else None |
| 219 | + |
| 220 | + |
| 221 | +def vector_extract_element(vec, position, *, loc=None, ip=None): |
| 222 | + """Extract one element of ``vec`` at ``position``, across CutDSL versions. |
| 223 | +
|
| 224 | + CutDSL renamed ``vector.extractelement`` to ``vector.extract`` in the 4.6 |
| 225 | + line. ``position`` is the Python element index; each branch builds the |
| 226 | + operand shape its binding actually expects: |
| 227 | +
|
| 228 | + - ``extractelement`` (4.5 line; also present in early 4.6): takes the |
| 229 | + index as a single i32 operand, so the constant is constructed here. |
| 230 | + - ``extract`` (4.6+): takes a sequence of index-typed dynamic operands |
| 231 | + plus a static-position array, so a constant position is ``extract(vec, |
| 232 | + [], [position], ...)``. |
| 233 | +
|
| 234 | + Preferring ``extractelement`` when both exist keeps the pre-4.6 code path |
| 235 | + byte-identical to what cuLA shipped before the gateway. |
| 236 | + """ |
| 237 | + vector_dialect = _load("vector") |
| 238 | + if _has_attribute_path(vector_dialect, ("extractelement",)): |
| 239 | + i32_ty = _load("ir").IntegerType.get_signless(32) |
| 240 | + index = _load("arith").constant(i32_ty, position, loc=loc, ip=ip) |
| 241 | + return vector_dialect.extractelement(vec, position=index, loc=loc, ip=ip) |
| 242 | + return vector_dialect.extract(vec, [], [position], loc=loc, ip=ip) |
0 commit comments