cutlass/examples/python/deprecated/00_basic_gemm.ipynb

Basic example of using the CUTLASS Python interface

This notebook walks through a basic example of using the CUTLASS Python interface to declare, compile, and run GEMMs.

Prerequisites for running on Colab

This notebook requires an NVIDIA GPU. If nvidia-smi fails, go to Runtime -> Change runtime type -> Hardware accelerator and confirm a GPU is selected.

In [ ]:

!#nvidia-smi

If running on Colab, you will need to install the CUTLASS Python interface. To do so, uncomment the following line and run the cell:

In [ ]:

!#pip install nvidia-cutlass

General setup

We first import various packages needed for the example and construct the input and output tensors that will be used in our example.

In [ ]:

import numpy as np
import random

import cutlass

# This controls whether the C++ GEMM declaration will be printed at each step. 
# Set to `False` to omit this information.
print_module = True

m = 128
n = m
k = m

dtype = np.float16
type_A = np.float16
type_B = np.float16
type_C = np.float16
type_D = np.float16

np.random.seed(1234)
random.seed(1234)
scope_min = -4
scope_max = 4
tensor_A = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(m, k)).astype(type_A))
tensor_B = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(k, n)).astype(type_B))
tensor_C = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(m, n)).astype(type_C))

alpha = np.float16(1.)
beta = np.float16(0.)

tensor_D = np.zeros(tensor_C.shape).astype(type_D)

Declaring and running a GEMM

To get started, one only needs to provide the tensors declared above to the cutlass.op.Gemm call. This sets up a default GEMM operation for the given device on which you are running.

Assuming that we are running on SM80, this default to using a GEMM that leverages FP16 Tensor Core operations.

Calling plan.run() will generate the CUTLASS C++ kernel in question, compile it, and run it on the tensors we previously passed in. By setting print_module to true, the C++ code that is emitted is printed.

In [ ]:

# We specify `element_accumulator` here so as to match the kernel run by NumPy below. However,
# specifying `element_accumulator` is not required if it is the same as `element`
plan = cutlass.Gemm(element=dtype, layout=cutlass.LayoutType.RowMajor, element_accumulator=np.float32)
plan.run(tensor_A, tensor_B, tensor_C, tensor_D, print_module=print_module)

There are many other ways to construct a plan from cutlass.op.Gemm (e.g., by specifiying they types and layouts of each operand, by providing representative tensors as inputs). For more details on these, see the documentation in the cutlass.op.Gemm constructor.

We then compare the output to running the GEMM using NumPy.

In [ ]:

tensor_D_numpy = (alpha * (tensor_A @ tensor_B)) + (beta * tensor_C)
np.testing.assert_array_equal(tensor_D, tensor_D_numpy)

Note that one could use the same kernel just declared for tensors provided by other frameworks beyond NumPy, such as PyTorch or CuPy.

Changing operation modes

By default, the CUTLASS Python interface will try to use Tensor Core operations whenever possible. If the configuration provided to cutlass.op.Gemm is not supported on Tensor Cores, the interface will fall back to using a SIMT kernel.

The operation mode currently in use can be returned via the plan.opclass property. In this case Tensor Core operations.

In [ ]:

print(plan.opclass)

Suppose that we don't want to use Tensor Cores for this GEMM. One can change to using CUTLASS's SIMT GEMMs by setting the plan's opclass field.

As is shown in the printed output, the emitted kernel uses template parameters that fit CUTLASS's SIMT GEMMs.

Also notice that, this time around, we provided tensor parameters to plan.run(). One is free to provide different parameters to plan.run() than were passed in at the initial call to cutlass.op.Gemm, provided that the passed-in tensors have the same data type and layout as those passed in on intialization.

In [ ]:

tensor_D_simt = np.zeros(tensor_C.shape).astype(type_D)
plan.opclass = cutlass.OpcodeClass.Simt
plan.run(tensor_A, tensor_B, tensor_C, tensor_D_simt, alpha, beta, print_module=print_module)

If we compare the output of the Tensor Core and SIMT GEMMs we just ran we see that they are equal.

In [ ]:

np.testing.assert_array_equal(tensor_D, tensor_D_simt)

Running cached kernels

You may have noticed that the plan.run() calls for the previous two kernels took some time to execute. This is because the kernel being emitted had not yet been compiled.

CUTLASS caches compiled binaries so that recompilation isn't necessary every time a kernel is run. For example, if we change modes back to using Tensor Cores and call plan.run() again (with a different set of tensor parameters), you'll find the call to return much faster.

In [ ]:

m = 2400
n = 3232
k = 4096

tensor_A = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(m, k)).astype(type_A))
tensor_B = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(k, n)).astype(type_B))
tensor_C = np.ceil(np.random.uniform(low=scope_min, high=scope_max, size=(m, n)).astype(type_C))
tensor_D = np.zeros(tensor_C.shape).astype(type_D)

alpha = np.float16(1.)
beta = np.float16(2.)

plan.opclass = cutlass.OpcodeClass.TensorOp
plan.run(tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, print_module=print_module)

Running non-default GEMMs

The previous examples showed how it is simple to get started running a default GEMM kernel in CUTLASS. But, what do you do if you want a bit more control over the parameters to the GEMM?

Under the hood, CUTLASS enumerates the different GEMM configuration parameters possible for this kernel from the CUTLASS profiler. The code below shows how one can access the tile descriptions for the kernels (e.g., cluster, threadblock, and warp shape).

In [ ]:

tiles = plan.tile_descriptions()
print('{} tile descriptions returned'.format(len(tiles)))
num_print = 10
print('First {} tile descriptions are:'.format(num_print))
for td in tiles[:num_print]:
    print(td)

Next, we'll pick one of these configurations at random and compile and run it.

In [ ]:

tiles = [td for td in tiles if td.threadblock_shape[0] >= 128]
idx = random.randint(0, len(tiles)-1)
td = tiles[idx]
print('Tile description {} is: {}'.format(idx, td))
plan.compile(td)
plan.run(tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, print_module=print_module)

One can also change the swizzling function used by the kernel. For example, one can modify the kernel to use the stream K feature of CUTLASS via:

In [ ]:

# Stream K is exposed through the threadblock swizzle method for pre-SM90 kernels,
# and via the tile_scheduler attribute of the TileDescription for post-SM90 kernels
if plan.cc < 90:
    plan.swizzling_functor = cutlass.swizzle.ThreadblockSwizzleStreamK
    plan.run(tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, print_module=print_module)
else:
    # Stream-K is currently only supported for warp-specialized cooperative kernels
    td.kernel_schedule = cutlass.KernelScheduleType.TmaWarpSpecializedCooperative
    td.epilogue_schedule = cutlass.EpilogueScheduleType.TmaWarpSpecializedCooperative
    td.tile_scheduler = cutlass.TileSchedulerType.StreamK

    plan.compile(td)
    plan.run(tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, print_module=print_module)

Handling errors

The CUTLASS Python interface attempts to catch runtime and compilation errors in Python so as to provide more understandable error messages.

Here's an example in which we try to use too many stages for a given GEMM kernel. Normally, this would result in a runtime error due to the GPU having insufficient shared memory to launch the kernel with 8 stages. The CUTLASS Python interface is able to detect this issue before compiling the kernel, and reports it back to the user. Uncomment and run the code below to see this error.

In [ ]:

# td = tiles[0]
# td.stages = 8
# plan.compile(td)

Specializations for other data types

Various CUTLASS kernels specialized for specific data types can also be run via the Python interface.

For example, the code below shows how to declare and run a GEMM using the 3xTF32 feature (see corresponding C++ example here).

In [ ]:

from cutlass.backend.utils.device import device_cc

# 3xTF32 requires SM80 or higher
if device_cc() >= 80:
    plan = cutlass.op.Gemm(element=np.float32, layout=cutlass.LayoutType.RowMajor)
    plan.math_operation = cutlass.MathOperation.multiply_add_fast_f32

    # Create input/output tensors in FP32
    A, B = [np.ones((128, 128)).astype(np.float32) for _ in range(2)]
    C, D = [np.zeros((128, 128)).astype(np.float32) for _ in range(2)]

    # Run the GEMM
    plan.run(A, B, C, D, print_module=print_module)

Additionally, one can run CUTLASS's FP8 GEMMs if using a frontend library capable of allocating and initializing FP8 tensors (e.g., PyTorch)

In [ ]:

try:
    import torch
except ImportError:
    print("PyTorch is not available. Skipping FP8 example")
    import sys; sys.exit(0)

if not hasattr(torch, "float8_e4m3fn"):
    print("Version of PyTorch does not have the float8_e4m3fn data type. Skipping FP8 example")
    import sys; sys.exit(0)

# FP8 is supported through the CUTLASS Python interface on SM90 and higher
if device_cc() >= 90:
    plan = cutlass.op.Gemm(element=torch.float8_e4m3fn, element_C=torch.float32, element_accumulator=torch.float32,
                        layout_A=cutlass.LayoutType.RowMajor, layout_B=cutlass.LayoutType.ColumnMajor,
                        layout_C=cutlass.LayoutType.ColumnMajor)

    # Create input/output tensors in FP8
    A, B = [torch.ones((128, 128)).to(torch.float8_e4m3fn).to("cuda") for _ in range(2)]
    C, D = [torch.zeros((128, 128)).to(torch.float8_e4m3fn).to("cuda") for _ in range(2)]

    # Run the GEMM
    plan.run(A, B, C, D, print_module=print_module)