Emitters - CUTLASS Python

Hide navigation sidebar
Hide table of contents sidebar
Back to top
Toggle table of contents sidebar

# Emitters

## Common

Common utilities for emitting CUTLASS kernels

## PyTorch

Utilities for generating source for building a PyTorch CUDA extension that using a CUTLASS kernel. If specified, the extension can be JIT compiled via PyTorch’s `cpp_extension.load` method.

Example usage with JIT compilation:

    plan = cutlass.op.Gemm(element=torch.float32, layout=cutlass.LayoutType.RowMajor)
    op = plan.construct()
    mod = cutlass.emit.pytorch(op, 'cutlass_gemm', 80, jit=True)

    # Generate inputs for the GEMM
    A, B, C = [torch.ones((512, 512)).to('cuda') for _ in range(3)]

    # Run the module
    D = mod.run(A, B, C)

Example usage without JIT compilation:

    plan = cutlass.op.Gemm(element=torch.float32, layout=cutlass.LayoutType.RowMajor)
    op = plan.construct()
    cutlass.emit.pytorch(op, 'cutlass_gemm', 80, jit=False, sourcedir='output')

After this call, the directory `output` contains `setup.py`, `cutlass_gemm.cpp`, and `cutlass_gemm_kernel.cu`. The module can be built from within `output` by running: `TORCH_CUDA_ARCH_LIST="8.0" python setup.py develop --user`.

The module can later be used in Python via:

    import torch
    import cutlass_gemm

    # Generate inputs for the GEMM
    A, B, C = [torch.ones((512, 512)).to('cuda') for _ in range(3)]

    # Run the module
    D = cutlass_gemm.run(A, B, C)

cutlass.emit.pytorch.pytorch(op, name, cc, jit=False, sourcedir='')

  Generates source for building a PyTorch CUDA module that leverages the CUTLASS kernel specified by `op`. If the `jit` parameter is set to true, the module is just-in-time compiled, loaded, and returned.

  The result of this method is files within `sourcedir` that can be used for building a PyTorch module.

  Parameters:

    - op – operation to emit in the module

    - name (str) – name of the module to generate

    - cc (int) – compute capability of the device the module should target

    - jit (bool) – whether the module should be just-in-time compiled

    - sourcedir (str) – directory to which generated source files should be written

  Returns:

    loaded PyTorch module (if `jit=True`) or None