def set_num_sms()

in deep_gemm/jit_kernels/utils.py [0:0]

• 4 lines of code
• 1 McCabe index (conditional complexity)

def set_num_sms(num_sms: int) -> None:
    """
    Set the maximum SM count for all GEMM kernels to use.

    Arguments:
        num_sms: the desired maximum SM count for all GEMM kernels to use.
    """
    global _num_sms
    assert 0 < num_sms <= torch.cuda.get_device_properties(device='cuda').multi_processor_count
    _num_sms = num_sms