2292557fd6
dist/include/ contains the MLX and MLX-C headers needed for CGo compilation. Without these, go-mlx cannot be used as a module dependency (headers not found in module cache). Libraries (dylib/metallib) are still gitignored — users build those locally via cmake. Co-Authored-By: Virgil <virgil@lethean.io>
46 lines
1.2 KiB
C++
46 lines
1.2 KiB
C++
// Copyright © 2025 Apple Inc.
|
|
|
|
// This file include utilities that are used by C++ code (i.e. .cpp files).
|
|
|
|
#pragma once
|
|
|
|
#include "mlx/array.h"
|
|
#include "mlx/backend/cuda/allocator.h"
|
|
#include "mlx/backend/cuda/cuda_utils.h"
|
|
|
|
namespace mlx::core {
|
|
|
|
template <typename T>
|
|
inline uint max_occupancy_block_dim(T kernel) {
|
|
int _, block_dim;
|
|
if constexpr (std::is_same_v<T, CUfunction>) {
|
|
CHECK_CUDA_ERROR(
|
|
cuOccupancyMaxPotentialBlockSize(&_, &block_dim, kernel, 0, 0, 0));
|
|
} else {
|
|
CHECK_CUDA_ERROR(
|
|
cudaOccupancyMaxPotentialBlockSize(&_, &block_dim, kernel));
|
|
}
|
|
return block_dim;
|
|
}
|
|
|
|
template <typename T>
|
|
inline T* gpu_ptr(array& arr) {
|
|
return reinterpret_cast<T*>(
|
|
static_cast<char*>(
|
|
static_cast<cu::CudaBuffer*>(arr.buffer().ptr())->data) +
|
|
arr.offset());
|
|
}
|
|
|
|
// For const array, keep constness in pointer unless it is untyped.
|
|
template <typename T>
|
|
inline std::conditional_t<std::is_same_v<T, void>, void*, const T*> gpu_ptr(
|
|
const array& arr) {
|
|
return gpu_ptr<T>(const_cast<array&>(arr));
|
|
}
|
|
|
|
struct Dtype;
|
|
|
|
// Convert Dtype to CUDA C++ types.
|
|
const char* dtype_to_cuda_type(const Dtype& dtype);
|
|
|
|
} // namespace mlx::core
|