Function device_mem_alloc

python/cutlass_cppgen/backend/memory_manager.py:102–109

(size)

def device_mem_alloc(size):
    if cutlass_cppgen.use_rmm:
        return rmm.DeviceBuffer(size=size)
    else:
        err, ptr = cudart.cudaMalloc(size)
        if err != cudart.cudaError_t.cudaSuccess:
            raise Exception(f"cudaMalloc failed with error {err}")
        return DevicePtrWrapper(ptr)


def align_size(size, alignment=256):