Operator interface

The Operator class is the central interface of Operator API. Each Operator wraps a single “kernel” written in DSL and exposes a uniform interface for executing it, generating its instances, and metadata that describes the kernel’s capabilities.

class cutlass.operators.base.Operator(metadata: OperatorMetadata)

Bases: ABC

Central class that provides a uniform interface over DSL kernels.

An Operator is Operator API’s central interface of instantiating, finding, and executing DSL kernels. A kernel is the raw DSL code that does the computation; an Operator wraps a kernel and exposes a uniform interface to it. Each Operator subclass fills in the kernel-agnostic interface with kernel-specific implementations for:

• correctly extracting arguments and compile/run the kernel

• generating instances of the Operator

• exposing the kernel’s properties via OperatorMetadata

Typically, DSL kernels are written as classes that are parameterized (e.g. by operand data types, tensor layouts, tile shapes, etc.) into a family of related kernel configurations, much like a “template”. Instantiating & compiling a kernel class with concrete parameters produces a single runnable kernel.

Likewise, an Operator subclass represents a family of parametrized kernels, wrapped in the uniform interface. Each instance of that subclass represents one fully-specified kernel configuration. The subclass is responsible for generating valid configurations of this family from the design space. All features that distinguish one Operator instance from another are exposed via a read-only metadata.

supported_args_type: ClassVar[type[RuntimeArguments]]

The subtype of RuntimeArguments that this Operator supports.

This is used by providers to filter operator classes before calling generate_operators(). Subclasses must override this, failing which an error is raised when the subclass is defined.

designed_for_min_cc: ClassVar[int]

Minimum compute capability that this operator class is designed for.

It is assumed that an Operator cannot run on an architecture with compute capability older than this. It indicates a broad class-level minimum, meant for rough filtering. However, it does not imply that the Operator can run on every newer architecture – operator.metadata.supported_targets provides finegrained ground-truth information about which architectures each Operator instance can run on.

property metadata: OperatorMetadata

The read-only metadata for the operator.

All features that distinguish one Operator from another (supported operand attributes, design choices, supported targets, optional epilogue, …) are exposed via metadata, and can be used for inspection and filtering.

final supports(

args: RuntimeArguments,

target_sm: TargetSm | str | None = None,

) → Status

Return whether the operator can be compiled or run with the provided arguments.

Parameters:

• args (RuntimeArguments) – Arguments with which the operator is to be compiled or run.

• target_sm (TargetSm | str | None) – An optional target compute capability to check support for.

Returns:

Status indicating whether the operator can be compiled or run with the provided arguments.

Return type:

Status

final compile(

args: RuntimeArguments,

target_sm: TargetSm | str | None = None,

) → CompiledArtifact

Compile the operator for the given arguments and target architecture.

This is implicitly called by run(), but can also be called explicitly to obtain a CompiledArtifact that and passed to run() later.

Parameters:

• args (RuntimeArguments) – Arguments to compile the operator with. The operator will be compiled to support the types and layouts of tensors in these args, according to any static information present in it (layouts, dtypes, etc.).

• target_sm (TargetSm | str | None) – Target compute capability to compile the operator for.

Returns:

The compiled artifact.

Return type:

CompiledArtifact

final run(

args: RuntimeArguments,

compiled_artifact: CompiledArtifact | None = None,

stream=None,

workspace: Workspace | None = None,

assume_supported_args: bool = False,

) → None

Execute the operator end-to-end with provided arguments.

Check that args are supported, compile if needed, create a default stream if needed, and launch the underlying kernel.

Parameters:

• args (RuntimeArguments) – Arguments with which to run the operator.

• compiled_artifact (CompiledArtifact | None) – Compiled artifact returned from the compile(). If None, the operator will first be compiled.

• stream (cuda.CUstream, torch.cuda.Stream, or other stream-like object, or None) – Stream to execute the operator on. If None, the default/null stream cuda.CUstream(0) is used.

• workspace (Workspace | None) – Allocation of workspace at least as large as the workspace size returned from the get_workspace_size(). If the operator does not require workspace, this can be None. If a workspace of inappropriate size is provided, the behavior is undefined and the operator may crash.

• assume_supported_args (bool) – By default, operator.supports(args) is called to check if the arguments are supported. If True, this check is skipped.

final get_workspace_size(

args: RuntimeArguments,

) → AllocationRequirement

Return the workspace requirement for the operator for the given arguments.

Parameters:

args (RuntimeArguments) – Arguments to the operator to get the workspace requirement for.

Returns:

Requirement describing the size & alignment of the workspace required by the operator.

Return type:

AllocationRequirement

final initialize_workspace(

args: RuntimeArguments,

workspace: Workspace,

) → None

Initialize the workspace for the operator.

Parameters:

• args (RuntimeArguments) – Arguments to initialize the workspace for.

• workspace (Workspace) – Workspace to initialize.

final classmethod generate_operators(

metadata_filter: Callable[[OperatorMetadata], bool],

epilogue_args: EpilogueArguments | None = None,

target_sm: TargetSm | str | None = None,

args: RuntimeArguments | None = None,

) → list[Operator]

Return a list of all supported operator configurations, optionally constrained to given parameters.

Parameters:

• metadata_filter (Callable[[OperatorMetadata], bool]) – A 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.

• epilogue_args (EpilogueArguments | None) – Optional arguments to pass to the operator’s epilogue.

• target_sm (TargetSm | str | None) – Optionally restrict generation to operators compatible with a given TargetSm.

• args (RuntimeArguments | None) – Optional arguments to filter operators by. If None, all operators are considered.

Returns:

List of all supported operator configurations.

Return type:

list[Operator]

Compiled artifacts

Operator.compile returns a CompiledArtifact that bundles the compiled kernel with the originating Operator and the target it was compiled for. It can be passed back to Operator.run to skip recompilation.

class cutlass.operators.artifact.CompiledArtifact(

compiled_obj: Any,

operator_obj: Operator,

compiled_for: TargetSm,

)

Bases: object

Lightweight wrapper over the result of compiling a DSL kernel.

For instance, when using CuTe DSL, it’s a wrapper over the result of cute.compile(). In addition, it stores which Operator and TargetSm an artifact was compiled for.

compiled_obj: Any

The result of compiling the Operator

operator_obj: Operator

The Operator object on which .compile() was called to generate this artifact

compiled_for: TargetSm

The target for which the Operator was compiled

Workspace

Some Operators need a scratch buffer to run. Operator.get_workspace_size returns an AllocationRequirement that the caller satisfies by passing a Workspace to Operator.run.

class cutlass.operators.workspace.Workspace(data: TensorLike | int, leading_zero_bytes: int = 0)

Bases: object

Container for a user-provided memory allocation used by an Operator.

A workspace is a user-provided, contiguous block of memory some Operators may use to store intermediate results, which may be used by the underlying kernel for any temporary storage, book keeping, intermediate results, semaphores etc. The workspace is typically allocated by the user and passed to the Operator via the Operator.run() method.

data: TensorLike | int

Tensor holding workspace data, or an integer address of the start of the workspace.

leading_zero_bytes: int = 0

Number of bytes from the beginning of the workspace have been set to zero by user.

class cutlass.operators.workspace.AllocationRequirement(size_bytes: int, ptr_alignment: int)

Bases: object

Describes requirements for a user-provided memory allocation.

size_bytes: int

The minimal size (in bytes) that must be allocated assuming perfect alignment.

ptr_alignment: int

Alignment requirement for the pointer to the allocation in bytes.

classmethod empty() → AllocationRequirement

Returns a requirement for an empty allocation.