CallConvVersion version;

inline bool operator== (const CallingConvention& other) const {

return version == other.version;

}

static PyTypeObject pytype;

PyObject* ret = CallingConvention::pytype.tp_new(&CallingConvention::pytype, nullptr, nullptr);

if (!ret) return {};

CallingConvention& cconv = py_unwrap<CallingConvention>(ret);

cconv.version = version;

return steal(ret);

if (s[0] == '1' && !s[1])

return get_cconv(CallConvVersion::CutilePython_V1);

return {};

explicit CudaLibrary(const DriverApi* driver, CUlibrary lib) : driver_(driver), lib_(lib) {}

CudaLibrary(CudaLibrary&& other) : driver_(other.driver_), lib_(other.lib_) {

other.lib_ = nullptr;

}

CudaLibrary(const CudaLibrary&) = delete;

void operator=(const CudaLibrary&) = delete;

~CudaLibrary() {

if (lib_) {

CUresult res = driver_->cuLibraryUnload(lib_);

CHECK(res == CUDA_SUCCESS);

}

}

const CUlibrary& get() const {

return lib_;

}

const DriverApi* driver_;

CUlibrary lib_;

CUlibrary lib;

CUresult res = driver->cuLibraryLoadData(&lib, code, nullptr, nullptr, 0,

nullptr, nullptr, 0);

if (res == CUDA_SUCCESS)

return CudaLibrary(driver, lib);

return raise(PyExc_RuntimeError, "Failed to load CUDA library: %s",

get_cuda_error(driver, res));

const char* cubin_data,

size_t cubin_size,

const char* func_name) {

(void) cubin_size;

Result<CudaLibrary> lib = load_cuda_library(driver, cubin_data);

if (!lib.is_ok()) return ErrorRaised;

CUkernel kernel;

CUresult res = driver->cuLibraryGetKernel(&kernel, lib->get(), func_name);

if (res == CUDA_SUCCESS)

return CudaKernel{std::move(*lib), kernel};

return raise(PyExc_RuntimeError, "Failed to get kernel %s from library: %s",

func_name, get_cuda_error(driver, res));

int* num_attrs, int* min_stack, int* stack_effect) {

FOREACH_SIZE_OPCODE(SIZE_OPCODE_PARSE);

return raise(PyExc_ValueError, "Invalid opcode string %R", opcode_str);

enum { kMaxStackDepth = 32 };

Vec<SizeOpcode> opcodes;

Vec<int64_t> op_attrs;

enum { kMaxRank = 5 };

CUtensorMapDataType dtype;

uint32_t item_size;

uint32_t rank;

uint32_t base_ptr_param_idx;

HostProgram shape_stride_program;

uint32_t box_dim[kMaxRank];

uint32_t traversal_steps[kMaxRank];

CUtensorMapInterleave interleave;

CUtensorMapSwizzle swizzle;

CUtensorMapL2promotion l2_promotion;

CUtensorMapFloatOOBfill oob_fill;

CudaKernel cukernel;

HostProgram dyn_smem_size_prog;

Vec<HoistedTensorMap> hoisted_tensor_maps;

ArenaOffset base_ptr_cuarg;

size_t length;

ArenaOffset item_offsets; // [length], each word contains an `ArenaOffset` points to the item

size_t item_size_words;

Vec<PyTypeObject*> pyarg_types;

Arena arena;

Vec<ArenaOffset> cuarg_offsets; // offsets into `arena`

Vec<void*> launch_params;

Vec<ListArg> list_args;

size_t total_list_data_size_words;

Vec<int64_t> constants;

CUcontext cuda_context;

LaunchHelper* next_free;

ArenaOffset offset = arena.size();

arena.resize(offset + count);

return offset;

static_assert(AlignmentBytes % sizeof(Word) == 0);

constexpr size_t AlignmentWords = AlignmentBytes / sizeof(Word);

size_t padded_size = ((arena.size() + AlignmentWords - 1) / AlignmentWords)

* AlignmentWords;

arena.resize(padded_size);

ArenaOffset offset = arena_alloc_words(helper.arena, 1);

helper.arena[offset] = word;

helper.cuarg_offsets.push_back(offset);

return offset;

void operator() (LaunchHelper* helper) const {

helper->next_free = g_helper_freelist;

g_helper_freelist = helper;

}

if (g_helper_freelist) {

LaunchHelper* ret = g_helper_freelist;

g_helper_freelist = ret->next_free;

return LaunchHelperPtr(ret);

} else {

return LaunchHelperPtr(new LaunchHelper());

}

Array,

Boolean,

Integer,

Float,

List,

// A torch.Tensor that we can access via torch._C._to_dlpack

TorchTensorDlpack,

// An object with __dlpack__ method

DlpackArray,

// An object with __cuda_array_interface__

CudaArray,

// Python `bool`,

PyBool,

// Python `int`,

PyLong,

// Python `float`

PyFloat,

// Python `list`

PyList

switch (k) {

case PythonArgKind::TorchTensorDlpack: return ParameterKind::Array;

case PythonArgKind::DlpackArray: return ParameterKind::Array;

case PythonArgKind::CudaArray: return ParameterKind::Array;

case PythonArgKind::PyBool: return ParameterKind::Boolean;

case PythonArgKind::PyLong: return ParameterKind::Integer;

case PythonArgKind::PyFloat: return ParameterKind::Float;

case PythonArgKind::PyList: return ParameterKind::List;

}

CHECK(false);

if (PyBool_Check(arg))

return PythonArgKind::PyBool;

if (PyLong_Check(arg))

return PythonArgKind::PyLong;

if (PyFloat_Check(arg))

return PythonArgKind::PyFloat;

if (PyList_Check(arg))

return PythonArgKind::PyList;

if (g_torch_Tensor_type && PyObject_TypeCheck(arg, g_torch_Tensor_type)) {

// Calling torch._C._to_dlpack(arg) is much faster than calling arg.__dlpack__()

// because it goes straight into C++ code, with no Python in between.

// So we always prefer that.

if (g_torch_to_dlpack_func)

return PythonArgKind::TorchTensorDlpack;

}

if (PyObject_HasAttr(arg, g___dlpack___pyunicode))

return PythonArgKind::DlpackArray;

if (PyObject_HasAttr(arg, g___cuda_array_interface___pyunicode))

return PythonArgKind::CudaArray;

return raise(PyExc_TypeError, "Unsupported argument type %s", Py_TYPE(arg)->tp_name);

unsigned xu = static_cast<unsigned char>(x);

unsigned yu = static_cast<unsigned char>(y);

return ((xu << 8) | yu);

if (!PyUnicode_Check(typestr)) {

PyErr_SetString(PyExc_TypeError, "__cuda_array_interface__['typestr'] is not a string");

return ErrorRaised;

}

Py_ssize_t len;

const char* str = PyUnicode_AsUTF8AndSize(typestr, &len);

if (!str) return ErrorRaised;

if (len < 3) {

PyErr_Format(PyExc_TypeError, "__cuda_array_interface__['typestr'] has invalid value %S",

typestr);

return ErrorRaised;

}

// TODO: support big endian one day?

if (str[0] != '<' && str[0] != '|') {

PyErr_SetString(PyExc_TypeError, "Only little-endian types are supported");

return ErrorRaised;

}

DLDataType ret;

ret.lanes = 1;

switch (str[1]) {

case 'b': ret.code = kDLBool; break;

case 'i': ret.code = kDLInt; break;

case 'u': ret.code = kDLUInt; break;

case 'f': ret.code = kDLFloat; break;

case 'V': ret.code = kDLBfloat; break;

case 'c': ret.code = kDLComplex; break;

default:

PyErr_Format(PyExc_TypeError, "Unsupported type code %c", str[1]);

return ErrorRaised;

}

// str[3] is safe to index because there is always a NUL byte at the end

switch (char_pair(str[2], str[3])) {

case char_pair('1', '\0'): ret.bits = 8; break;

case char_pair('2', '\0'): ret.bits = 16; break;

case char_pair('4', '\0'): ret.bits = 32; break;

case char_pair('8', '\0'): ret.bits = 64; break;

case char_pair('1', '6'):

if (!str[4]) {

ret.bits = 64;

break;

}

[[fallthrough]];

default:

PyErr_Format(PyExc_TypeError, "Unsupported byte size in typestr: %s", str + 2);

return ErrorRaised;

}

return ret;

DLDataType dtype;

size_t ndim;

unsigned index_bitwidth;

return static_cast<uint32_t>(dtype.code)

| (static_cast<uint32_t>(dtype.bits) << 8)

| (static_cast<uint32_t>(dtype.lanes) << 16);

return DLDataType{

.code = static_cast<uint8_t>(u & 0xff),

.bits = static_cast<uint8_t>((u >> 8) & 0xff),

.lanes = static_cast<uint16_t>((u >> 16) & 0xffff),

};

if (dtype.lanes != 1)

return raise(PyExc_TypeError, "Array dtypes with multiple lanes are not supported");

switch (u8_pair(dtype.code, dtype.bits)) {

case u8_pair(kDLBool, 8): return "bool_";

case u8_pair(kDLInt, 8): return "int8";

case u8_pair(kDLInt, 16): return "int16";

case u8_pair(kDLInt, 32): return "int32";

case u8_pair(kDLInt, 64): return "int64";

case u8_pair(kDLUInt, 8): return "uint8";

case u8_pair(kDLUInt, 16): return "uint16";

case u8_pair(kDLUInt, 32): return "uint32";

case u8_pair(kDLUInt, 64): return "uint64";

case u8_pair(kDLFloat, 16): return "float16";

case u8_pair(kDLFloat, 32): return "float32";

case u8_pair(kDLFloat, 64): return "float64";

case u8_pair(kDLBfloat, 16): return "bfloat16";

case u8_pair(kDLFloat8_e4m3fn, 8): return "float8_e4m3fn";

case u8_pair(kDLFloat8_e5m2, 8): return "float8_e5m2";

case u8_pair(kDLFloat8_e8m0fnu, 8): return "float8_e8m0fnu";

default:

return raise(PyExc_TypeError, "Unsupported array dtype");

}

PyObject* dtype_module = get_datatype_module();

if (!dtype_module) return {};

Result<const char*> name = dtype_name(dtype);

if (!name.is_ok()) return {};

return getattr(dtype_module, *name);

uint64_t dtype_u = static_cast<uint64_t>(dtype_as_uint(a.dtype));

uint64_t ndim_u = static_cast<uint64_t>(a.ndim);

uint64_t ibw_bit = (a.index_bitwidth == 64) ? 1ULL : 0ULL;

return static_cast<int64_t>(dtype_u | (ndim_u << 32) | (ibw_bit << 63));

uint64_t u = c;

uint32_t dtype = u & 0xffffffff;

uint32_t ndim = (u >> 32) & 0x7fffffff;

unsigned ibw = ((u >> 63) & 1) ? 64 : 32;

return {dtype_from_uint(dtype), ndim, ibw};

if (ndim == 0) return OK;

Word* shape = repr + 1 + ndim;

Word* stride = shape + ndim;

uint64_t prev_stride = 1;

(--stride)->i64 = 1;

for (size_t i = 0; i < ndim - 1; ++i) {

uint64_t new_stride = prev_stride * static_cast<uint64_t>((--shape)->i64);

if (index_bitwidth != 64 && new_stride > INT32_MAX)

return raise(PyExc_OverflowError, "stride is too big");

(--stride)->i64 = new_stride;

prev_stride = new_stride;

}

return OK;

const Word* array_repr, size_t ndim, unsigned dtype_bitwidth, unsigned index_bitwidth) {

ArraySpecializationBits ret = {};

void* data_ptr = array_repr[0].device_ptr;

const Word* shape_words = array_repr + 1;

const Word* stride_words = shape_words + ndim;

// Only specialize stride divisibility, stride 1 and shape divisibility for ndim <= TMA_MAX_NDIM

if (ndim <= TMA_MAX_NDIM) {

for (size_t i = 0; i < ndim; ++i) {

int64_t stride = stride_words[i].i64;

int64_t shape = shape_words[i].i64;

int64_t stride_bitwidth = stride * dtype_bitwidth;

int64_t shape_bitwidth = shape * dtype_bitwidth;

bool is_stride_byte_aligned = stride_bitwidth % BYTE_BITWIDTH == 0;

bool is_stride_16_byte_divisible =

(stride_bitwidth / BYTE_BITWIDTH) % DIVISOR_16 == 0;

bool is_shape_byte_aligned = shape_bitwidth % BYTE_BITWIDTH == 0;

bool is_shape_divisible_by_16 = shape % DIVISOR_16 == 0;

if (is_stride_byte_aligned && is_stride_16_byte_divisible)

ret.stride_16byte_divisible |= 1u << i;

if (stride == 1)

ret.stride_one |= 1u << i;

if (is_shape_byte_aligned && is_shape_divisible_by_16)

ret.shape_divisible_by_16 |= 1u << i;

}

}

// extract base pointer divisibility

intptr_t data_ptr_int = reinterpret_cast<intptr_t>(data_ptr);

ret.baseptr_16byte_aligned = data_ptr_int % DIVISOR_16 == 0;

// check elements disjoint.

// sort by stride. the smallest stride indicates the contiguous axis

// of the underlying array.

Vec<std::pair<int64_t, int64_t>> strides_and_shape(ndim);

for (size_t i = 0; i < ndim; ++i) {

strides_and_shape[i] = {stride_words[i].i64, shape_words[i].i64};

}

std::sort(strides_and_shape.begin(), strides_and_shape.end());

// disjointness check:

// - 0 dimension array elements are always disjoint.

// - >0 dimension array elements are disjoint if every stride is positive

// and greater than or equal to the product of the previous stride and

// the previous shape.

bool elems_disjoint = (ndim == 0) || (strides_and_shape[0].first > 0);

for (size_t i = 0; i + 1 < ndim; ++i) {

int64_t prev_stride = strides_and_shape[i].first;

int64_t prev_shape = strides_and_shape[i].second;

int64_t cur_stride = strides_and_shape[i + 1].first;

elems_disjoint &= (

cur_stride > 0 && cur_stride >= prev_stride * prev_shape);

}

ret.disjoint_elements = elems_disjoint;

return ret;

const int64_t* data;

size_t len;

int64_t next() {

CHECK(len);

int64_t ret = *data;

++data, --len;

return ret;

}

void* device_ptr = nullptr;

uint64_t bits = -1;

void update(const Arena& arena, const ArrayRepr& ar) {

const Word* repr = arena.data() + ar.repr;

device_ptr = repr[0].device_ptr;

bits &= compute_array_specialization_bits(

repr, ar.arrty.ndim, ar.arrty.dtype.bits * ar.arrty.dtype.lanes,

ar.arrty.index_bitwidth).u64;

}

void finalize(const DriverApi* driver, const ArrayType& arrty, LaunchHelper& helper) {

helper.constants.push_back(pack_array_type(arrty));

if (!helper.cuda_context) {

driver->cuPointerGetAttribute(&helper.cuda_context, CU_POINTER_ATTRIBUTE_CONTEXT,

reinterpret_cast<CUdeviceptr>(device_ptr));

}

helper.constants.push_back(bits);

}

ArrayType arrty = unpack_array_type(cursor.next());

ArraySpecializationBits special_bits;

special_bits.u64 = cursor.next();

unsigned index_bitwidth = arrty.index_bitwidth;

// Only int32 or int64 are supported now.

CHECK(index_bitwidth == 32 || index_bitwidth == 64);

PyObject* signature_module = get_signature_module();

if (!signature_module) return {};

PyPtr constraint_class = getattr(signature_module, "ArrayConstraint");

if (!constraint_class) return {};

PyPtr args = steal(PyTuple_New(0));

if (!args) return {};

PyPtr dtype = dtype_to_python(arrty.dtype);

if (!dtype) return {};

PyPtr index_dtype = dtype_to_python(

DLDataType{kDLInt, static_cast<uint8_t>(index_bitwidth), 1});

if (!index_dtype) return {};

PyPtr constant_strides = steal(PyTuple_New(arrty.ndim));

if (!constant_strides) return {};

PyPtr stride_divisible_by = steal(PyTuple_New(arrty.ndim));

if (!stride_divisible_by) return {};

PyPtr shape_divisible_by = steal(PyTuple_New(arrty.ndim));

if (!shape_divisible_by) return {};

PyPtr zero = steal(PyLong_FromLong(0));

if (!zero) return {};

PyPtr one = steal(PyLong_FromLong(1));

if (!one) return {};

PyPtr sixteen = steal(PyLong_FromLong(DIVISOR_16));

if (!sixteen) return {};

PyPtr stride_divisor = one;

constexpr unsigned divisor16_bits = DIVISOR_16 * BYTE_BITWIDTH;

if (divisor16_bits % arrty.dtype.bits == 0) {

stride_divisor = steal(PyLong_FromLong(divisor16_bits / arrty.dtype.bits));

if (!stride_divisor) return {};

}

for (size_t i = 0; i < arrty.ndim; ++i) {

PyObject* obj = special_bits.is_stride_one(i) ? one.get() : Py_None;

PyTuple_SET_ITEM(constant_strides.get(), i, Py_NewRef(obj));

obj = special_bits.is_stride_16byte_divisible(i) ? stride_divisor.get() : one.get();

PyTuple_SET_ITEM(stride_divisible_by.get(), i, Py_NewRef(obj));

obj = special_bits.is_shape_divisible_by_16(i) ? sixteen.get() : one.get();

PyTuple_SET_ITEM(shape_divisible_by.get(), i, Py_NewRef(obj));

}

PyPtr kwargs = steal(Py_BuildValue(

"{sO sI sO sO sO s() sO sO sO sO}",

"dtype", dtype.get(),

"ndim", static_cast<unsigned>(arrty.ndim),

"index_dtype", index_dtype.get(),

"stride_constant", constant_strides.get(),

"stride_lower_bound_incl", zero.get(),

"alias_groups",

"may_alias_internally", special_bits.disjoint_elements ? Py_False : Py_True,

"stride_divisible_by", stride_divisible_by.get(),

"shape_divisible_by", shape_divisible_by.get(),

"base_addr_divisible_by",

special_bits.baseptr_16byte_aligned ? sixteen.get() : one.get()));

if (!kwargs) return {};

return steal(PyObject_Call(constraint_class.get(), args.get(), kwargs.get()));

Arena& arena) {

PyPtr dict = steal(PyObject_GetAttr(pyobj, g___cuda_array_interface___pyunicode));

if (!PyDict_Check(dict.get())) {

PyErr_SetString(PyExc_TypeError,

"__cuda_array_interface__ returned a non-dictionary object");

return ErrorRaised;

}

UNPACK_ARRAY_INTERFACE(dict, typestr);

UNPACK_ARRAY_INTERFACE(dict, shape);

UNPACK_ARRAY_INTERFACE(dict, data);

// Parse the dtype

Result<DLDataType> dtype = parse_typestr(typestr);

if (!dtype.is_ok()) return ErrorRaised;

// Parse the data pointer

if (!PyTuple_Check(data) || PyTuple_GET_SIZE(data) != 2) {

PyErr_SetString(PyExc_TypeError,

"__cuda_array_interface['data'] is not a tuple of length 2");

return ErrorRaised;

}

PyObject* data_ptr_pylong = PyTuple_GET_ITEM(data, 0);

if (!PyLong_Check(data_ptr_pylong)) {

PyErr_SetString(PyExc_TypeError, "__cuda_array_interface['data'][0] is not an integer");

return ErrorRaised;

}

intptr_t data_ptr_int = pylong_as<intptr_t>(data_ptr_pylong);

if (PyErr_Occurred()) return ErrorRaised;

Py_ssize_t ndim = PyTuple_GET_SIZE(shape);

ASSERT_NDIM(ndim);

ArenaOffset repr_offset = arena_alloc_words(arena, 1 + 2 * ndim);

arena[repr_offset].device_ptr = reinterpret_cast<void*>(data_ptr_int);

// Parse the shape

if (!PyTuple_Check(shape))

return raise(PyExc_TypeError, "__cuda_array_interface['shape'] is not a tuple");

for (Py_ssize_t i = 0; i < ndim; ++i) {

int64_t size = pylong_as<int64_t>(PyTuple_GET_ITEM(shape, i));

if (PyErr_Occurred()) return ErrorRaised;

arena[repr_offset + 1 + i].i64 = size;

}

// Parse the strides

PyObject* strides = PyDict_GetItem(dict.get(), g_strides_pyunicode);

if (PyErr_Occurred()) return ErrorRaised;

if (!strides || strides == Py_None) {

if (!fill_row_major_strides(index_bitwidth, arena.data() + repr_offset, ndim))

return ErrorRaised;

} else if (PyTuple_Check(strides)) {

// Only byte-aligned types should be supported by __cuda_array_interface__

uint8_t dtype_bytewidth = dtype->bits / BYTE_BITWIDTH;

for (Py_ssize_t i = 0; i < ndim; ++i) {

int64_t stride = pylong_as<int64_t>(PyTuple_GET_ITEM(strides, i));

if (PyErr_Occurred()) return ErrorRaised;

arena[repr_offset + 1 + ndim + i].i64 = static_cast<int64_t>(

stride / dtype_bytewidth);

}

} else {

return raise(PyExc_TypeError, "__cuda_array_interface['strides'] can only be"

" absent, None, or a tuple");

}

return ArrayRepr {

.arrty = {

.dtype = *dtype,

.ndim = static_cast<size_t>(ndim),

.index_bitwidth = index_bitwidth

},

.repr = repr_offset

};

Arena& arena) {

void* ptr = PyCapsule_GetPointer(dlpack_capsule, "dltensor");

if (!ptr) return ErrorRaised;

DLManagedTensor* tensor = static_cast<DLManagedTensor*>(ptr);

if (tensor->dl_tensor.device.device_type != kDLCUDA)

return raise(PyExc_ValueError, "Input array is not on a CUDA device");

// TODO: check device ID

void* data_ptr = static_cast<char*>(tensor->dl_tensor.data) + tensor->dl_tensor.byte_offset;

uint32_t ndim = tensor->dl_tensor.ndim;

ASSERT_NDIM(ndim);

ArenaOffset repr_offset = arena_alloc_words(arena, 1 + 2 * ndim);

arena[repr_offset].device_ptr = data_ptr;

for (uint32_t i = 0; i < ndim; ++i) {

if (index_bitwidth != 64 && (tensor->dl_tensor.shape[i] < INT32_MIN

|| tensor->dl_tensor.shape[i] > INT32_MAX))

return raise(PyExc_OverflowError, "shape is too big");

arena[repr_offset + 1 + i].i64 = tensor->dl_tensor.shape[i];

}

if (!tensor->dl_tensor.strides) {

if (!fill_row_major_strides(index_bitwidth, arena.data() + repr_offset, ndim))

return ErrorRaised;

} else {

for (uint32_t i = 0; i < ndim; ++i) {

if(index_bitwidth != 64 && (tensor->dl_tensor.strides[i] < INT32_MIN

|| tensor->dl_tensor.strides[i] > INT32_MAX))

return raise(PyExc_OverflowError, "stride is too big");

arena[repr_offset + 1 + ndim + i].i64 = tensor->dl_tensor.strides[i];

}

}

ArrayRepr ret = {

.arrty = {

.dtype = tensor->dl_tensor.dtype,

.ndim = ndim,

.index_bitwidth = index_bitwidth

},

.repr = repr_offset

};