nvidia-cutlass-operators 0.1.0__py3-none-any.whl

cutlass/operators/__init__.py

@@ -0,0 +1,136 @@
+# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from cutlass.operators import mma, workspace
+from cutlass.operators.arch import ArchPortability, TargetSm
+from cutlass.operators.arguments import (
+    DenseTensor,
+    EpilogueArguments,
+    GemmArguments,
+    GroupedGemmArguments,
+    Operand,
+    PerformanceControls,
+    RuntimeArguments,
+    ScaledOperand,
+    ScaleMode,
+    ScaleSwizzleMode,
+)
+from cutlass.operators.artifact import CompiledArtifact
+from cutlass.operators.base import Operator
+from cutlass.operators.config import GlobalOptions
+from cutlass.operators.manifest import Manifest
+from cutlass.operators.metadata import OperatorMetadata
+from cutlass.operators.providers import (
+    CuTeDSLProvider,
+    Provider,
+    available_providers,
+    register_provider,
+)
+from cutlass.operators.status import Status
+from cutlass.operators.typing import NumericLike, TensorLike
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+__version__ = "0.1.0"
+
+
+def get_operators(
+    args: RuntimeArguments | None = None,
+    metadata_filter: Callable[[OperatorMetadata], bool] | None = None,
+    target_sm: TargetSm | str | None = None,
+    providers: list[Provider] | None = None,
+) -> list[Operator]:
+    """Return Operators that match the given arguments, metadata filter, and target.
+
+    Args:
+        args (RuntimeArguments | None): Runtime arguments describing the
+            operator invocation (e.g. :class:`GemmArguments`). When ``None``,
+            no argument-based filtering is applied.
+        metadata_filter (Callable[[OperatorMetadata], bool] | None): An
+            optional Callable that takes OperatorMetadata as input and returns
+            a boolean indicating if it should be considered for inclusion in
+            results. The result is an intersection of operators filtered by the
+            callable and by other parameters passed to this method.
+        target_sm (TargetSm | str | None): Compute capability to target (e.g.
+            ``"100a"`` or a :class:`TargetSm` instance). Filters Operators that
+            cannot run on this target.
+        providers (list[Provider] | None): Optional list of Providers to
+            restrict discovery to (e.g. ``[ops.CuTeDSLProvider]``).
+
+    Returns:
+        list[Operator]: Operators matching all filters.
+    """
+    return Manifest.get_operators(args, metadata_filter, target_sm, providers)
+
+
+__all__ = [
+    # Central class exposing the Operator interface
+    "Operator",
+    # Runtime arguments
+    "RuntimeArguments",
+    "PerformanceControls",
+    "EpilogueArguments",
+    "GemmArguments",
+    "GroupedGemmArguments",
+    # Operands
+    "Operand",
+    "DenseTensor",
+    "ScaledOperand",
+    "ScaleMode",
+    "ScaleSwizzleMode",
+    # Operator Discovery
+    "get_operators",
+    "Manifest",
+    "OperatorMetadata",
+    # Misc. core types
+    "Status",
+    "CompiledArtifact",
+    # Arch
+    "TargetSm",
+    "ArchPortability",
+    # Typing markers
+    "TensorLike",
+    "NumericLike",
+    # Configuration
+    "GlobalOptions",
+    # Provider management
+    "Provider",  # base class for all Providers
+    "CuTeDSLProvider",  # CuTeDSLProvider holding registry of kernels written in CuTe DSL
+    "available_providers",  # global list of available, registered Providers
+    "register_provider",  # decorator to register a new Provider
+    # Submodules to access less commonly used public surfaces
+    "mma",
+    "workspace",
+]

cutlass/operators/arch.py

@@ -0,0 +1,219 @@
+# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+from __future__ import annotations
+
+import enum
+from dataclasses import dataclass
+from enum import auto as enum_auto
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from cutlass.operators.metadata import DesignMetadata, OperandsMetadata
+
+
+@dataclass(init=False)
+class TargetSm:
+    """Target compute capability & portability to compile for."""
+
+    cc: int
+    portability: ArchPortability
+
+    def __init__(
+        self,
+        sm_str: str | None = None,
+        *,
+        cc: int | None = None,
+        portability: ArchPortability | None = None,
+    ):
+        """Create a ``TargetSm``.
+
+        Args:
+            sm_str (str | None): String like "100", "100a", or "100f".
+            cc (int | None): Compute capability as integer (e.g., 100).
+            portability (ArchPortability | None): Architecture portability. If not
+                provided, defaults to ArchConditional if cc >= 90, else Portable.
+
+        Note:
+            Must use either sm_str or cc/portability, not both.
+
+        Example:
+            >>> TargetSm("100a")
+            >>> TargetSm(cc=100)
+            >>> TargetSm(cc=100, portability=ArchPortability.ArchConditional)
+
+        Raises:
+            ValueError: If both sm_str and cc/portability are provided, or neither.
+            ValueError: If portability is not valid for the given cc.
+        """
+        if sm_str is not None:
+            if cc is not None or portability is not None:
+                raise ValueError("Cannot specify sm_str with cc/portability")
+            self.cc, self.portability = self._parse_str(sm_str)
+        elif cc is not None:
+            self.cc = cc
+            self.portability = portability or (
+                ArchPortability.ArchConditional
+                if cc >= 90
+                else ArchPortability.Portable
+            )
+        else:
+            raise ValueError("Must specify either sm_str or cc/portability")
+
+        if self.portability == ArchPortability.ArchConditional and self.cc < 90:
+            raise ValueError(
+                f"TargetSm {self.__str__()} is invalid. ArchConditional targets must be Hopper (sm90) or newer"
+            )
+        if self.portability == ArchPortability.FamilyPortable and self.cc < 100:
+            raise ValueError(
+                f"TargetSm {self.__str__()} is invalid. FamilyPortable targets must be Blackwell (sm100) or newer"
+            )
+
+    @classmethod
+    def ensure(cls, value: TargetSm | str | None) -> TargetSm | None:
+        """Coerce a string to ``TargetSm``, or return as-is if already one."""
+        if value is None or isinstance(value, cls):
+            return value
+        if isinstance(value, str):
+            return cls(value)
+        raise TypeError(f"Expected TargetSm or str. Got {value} of type {type(value)}")
+
+    @staticmethod
+    def get_supported_targets(
+        design: DesignMetadata, operands: OperandsMetadata
+    ) -> list[TargetSm]:
+        """Get the supported targets for a given design and operands.
+
+        Raises:
+            ValueError: If ``design`` does not have a ``mma_instruction_type`` field.
+        """
+        # For now, just return the supported instruction.
+        # It can be extended to any other rules that are introduced.
+        metadata_to_check = []
+
+        if not hasattr(design, "mma_instruction_type"):
+            raise ValueError(
+                "Cannot determine supported targets with missing mma_instruction_type field in "
+                f"design metadata: {design}"
+            )
+        mma_supported_targets = design.mma_instruction_type.supported_targets(
+            design, operands
+        )
+
+        metadata_to_check.append(mma_supported_targets)
+
+        # check against tile scheduler supported targets
+        if hasattr(design, "tile_scheduler") and design.tile_scheduler:
+            tile_scheduler_supported_targets = design.tile_scheduler.supported_targets(
+                design, operands
+            )
+            metadata_to_check.append(tile_scheduler_supported_targets)
+        return list(set.intersection(*map(set, metadata_to_check)))
+
+    def is_portable_to(self, other: TargetSm | str) -> bool:
+        """Check if this target can compile/run on architecture described by `other` TargetSm."""
+        other = TargetSm.ensure(other)
+        match self.portability:
+            case ArchPortability.Portable:
+                return self.cc <= other.cc
+            case ArchPortability.FamilyPortable:
+                return self.major == other.major and self.cc <= other.cc
+            case ArchPortability.ArchConditional:
+                return self == other
+            case _:
+                raise NotImplementedError
+
+    def supports_operators_from(
+        self, operator_targets: Iterable[TargetSm | str]
+    ) -> bool:
+        """Check if this target can compile/run operators designed for any of the given targets."""
+        return any(
+            TargetSm.ensure(target).is_portable_to(self) for target in operator_targets
+        )
+
+    @property
+    def major(self) -> int:
+        """Get the major version of the target SM."""
+        return self.cc // 10
+
+    @property
+    def minor(self) -> int:
+        """Get the minor version of the target SM."""
+        return self.cc % 10
+
+    @staticmethod
+    def _parse_str(s: str) -> tuple[int, ArchPortability]:
+        """Parse a compile target string (e.g., "100", "100a", "100f")."""
+        portability_map = {
+            "a": ArchPortability.ArchConditional,
+            "f": ArchPortability.FamilyPortable,
+        }
+
+        suffix = s[-1]
+        try:
+            if suffix.isdigit():
+                return (int(s), ArchPortability.Portable)
+            return (int(s[:-1]), portability_map[suffix])
+        except (ValueError, KeyError):
+            raise ValueError(
+                f"Invalid TargetSm: '{s}'. Expected form: '100', '100a', '100f'"
+            ) from None
+
+    def __str__(self) -> str:
+        suffix = (
+            self.portability.name[0].lower()
+            if self.portability != ArchPortability.Portable
+            else ""
+        )
+        return f"{self.cc}{suffix}"
+
+    def __hash__(self) -> int:
+        return hash((self.cc, self.portability))
+
+
+class ArchPortability(enum.Enum):
+    """Portability of a compiled Operator to other architectures/compute capabilities.
+
+    See:
+        https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/compute-capabilities.html#feature-set-compiler-targets
+    """
+
+    Portable = enum_auto()
+    """Portable to future architectures (e.g. sm_100 without "a" or "f")."""
+
+    FamilyPortable = enum_auto()
+    """Portable to future architectures within the same family (e.g. sm_100f).
+    Only applicable to Blackwell and newer architectures (cc >= 100).
+    """
+
+    ArchConditional = enum_auto()
+    """Not portable to any other architecture (e.g. sm_100a).
+    Only applicable to Hopper and newer architectures (cc >= 90).
+    """

cutlass/operators/arguments/__init__.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
+from .base import (
+    Operand,
+    PerformanceControls,
+    RuntimeArguments,
+)
+from .epilogue import EpilogueArguments
+from .gemm import GemmArguments, GemmProblemSize
+from .grouped_gemm import GroupedGemmArguments
+from .operand import (
+    DenseTensor,
+    ScaledOperand,
+    ScaleMode,
+    ScaleSwizzleMode,
+)
+
+__all__ = [
+    # Top-level arguments & constituents
+    "RuntimeArguments",
+    "PerformanceControls",
+    "EpilogueArguments",
+    # Various RuntimeArguments subtypes
+    "GemmArguments",
+    "GemmProblemSize",
+    "GroupedGemmArguments",
+    # Operands & related classes
+    "Operand",
+    "DenseTensor",
+    "ScaledOperand",
+    "ScaleMode",
+    "ScaleSwizzleMode",
+]

cutlass/operators/arguments/base.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+from __future__ import annotations
+
+import copy
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field, fields
+from typing import Any, final, get_type_hints
+
+from cutlass.operators.typing import NumericLike, TensorLike
+from cutlass.operators.utils.dtype import to_cutlass_type
+from cutlass.operators.utils.tensor import TensorWrapper
+
+
+@dataclass
+class RuntimeArguments:
+    """Describes the operands and all other arguments passed to an Operator at runtime.
+
+    It contains runtime operands (usually tensors) passed to the operation, as well as
+    any custom epilogue fusions, and runtime performance controls.
+
+    This is an abstract base class, whose subclass describes the operation type itself
+    (e.g. GemmArguments, GroupedGemmArguments, etc.). All operators that implement the same
+    operation type accept the same RuntimeArguments subclass.
+    """
+
+    performance: PerformanceControls | None = field(default=None, kw_only=True)
+    """Optional runtime performance controls passed to the Operator"""
+
+    def _validate(self):
+        """Checks that the arguments are valid.
+
+        This is run before all fields have been converted to TensorWrapper and cutlass.Numeric.
+        """
+
+    def __post_init__(self):
+        _convert_to_internal_types(self)
+
+
+@dataclass
+class PerformanceControls:
+    """Optional runtime performance controls passed to the Operator.
+
+    Some operators may support performance options that can be controlled at runtime.
+    This class is the general container for all such controls.
+    """
+
+
+class Operand(ABC):
+    """Base class for all operands to Operators, which encapsulates one or more TensorLike objects.
+
+    In the most basic case, an Operand enacapsulates a single tensor.
+
+    In more complex cases, an Operand may encapsulate multiple tensors that encapsulate a single logical operand.
+    For instance, a :class:`~cutlass.operators.ScaledOperand` encapsulates a quantized and scale tensor, that together
+    reconstruct the logical value of the operand.
+    """
+
+    @final
+    def copy(self) -> Operand:
+        """Returns a copy of the operand. Does not copy the underlying tensor."""
+        return self.__copy__()
+
+    @abstractmethod
+    def __copy__(self) -> Operand:
+        """Returns a copy of the operand. Does not copy the underlying tensor."""
+        raise NotImplementedError
+
+    def _convert_to_internal_types(self, metadata: dict[str, Any] = None):
+        _convert_to_internal_types(self, metadata=metadata)
+
+
+def _convert_to_internal_types(caller, metadata: dict[str, Any] = None):
+    """Converts fields of the caller to internal types.
+
+    Current fields that are converted:
+    * ``TensorLike`` -> ``TensorWrapper``
+    * ``NumericLike`` -> ``cutlass.Numeric``
+    * Classes that implement ``_convert_to_internal_types`` -> their internal types
+
+    Args:
+        caller (Any): The caller object to convert the fields of
+        metadata (dict[str, Any] | None): Additional metadata to be used for conversion
+    """
+    type_hints = get_type_hints(type(caller))
+    for f in fields(caller):
+        hint = type_hints.get(f.name)
+        value = getattr(caller, f.name)
+
+        global_metadata = {} if metadata is None else copy.deepcopy(metadata)
+        global_metadata.update(f.metadata)
+
+        if isinstance(value, TensorWrapper):
+            # No conversion needed
+            setattr(caller, f.name, value)
+        elif hint is TensorLike:
+            # Find all fields that are annotated as TensorLike,
+            # and wrap them in TensorWrapper
+            setattr(caller, f.name, TensorWrapper(value, **global_metadata))
+        elif hint is NumericLike:
+            # Find all fields that are annotated as NumericLike,
+            # and convert them to cutlass.Numeric
+            setattr(caller, f.name, to_cutlass_type(value))
+        elif hasattr(value, "_convert_to_internal_types"):
+            # If the field is an instance of a class that implements
+            # _convert_to_internal_types, convert it to internal types
+            value._convert_to_internal_types(metadata=global_metadata)
+            setattr(caller, f.name, value)