L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix some MSVC compilation errors (#3048)

Browse Source
This commit is contained in:
Cheng
2026-01-24 07:56:56 +09:00
committed by GitHub
parent fd27829efa
commit 8e93b7448c
18 changed files with 52 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx/backend/cuda/allocator.cpp
+1 -1
View File
@@ -3,11 +3,11 @@
#include "mlx/backend/cuda/allocator.h"
#include "mlx/backend/cuda/device.h"
#include "mlx/backend/cuda/utils.h"
#include "mlx/memory.h"
#include "mlx/utils.h"
#include <cuda_runtime.h>
#include <fmt/format.h>
#include <unistd.h>
#include <cassert>
mlx/backend/cuda/compiled.cpp
+1 -1
View File
@@ -258,7 +258,7 @@ void Compiled::eval_gpu(
"mlx::core::cu::{}_contiguous<int64_t, {}>",
lib_name(),
work_per_thread));
for (auto wpt : std::array<int, 2>{1, work_per_thread}) {
for (int wpt : {1, work_per_thread}) {
for (int i = 1; i <= MAX_NDIM; ++i) {
kernel_names.push_back(fmt::format(
"mlx::core::cu::{}_strided<{}, uint32_t, {}>", lib_name(), i, wpt));
mlx/backend/cuda/cudnn_utils.h
+3 -3
View File
@@ -72,13 +72,13 @@ inline fe::DataType_t dtype_to_cudnn_type(Dtype dtype) {
// There are 2 differences from the const_param util from kernel_utils.cuh:
// 1. The rest of array is filled with 0.
// 2. This util can be used in .cpp files.
template <int NDIM = MAX_NDIM, typename T, template <typename U> class Vec>
inline std::array<T, NDIM> vector_key(const Vec<T>& vec) {
template <int NDIM = MAX_NDIM, typename Vec>
inline std::array<typename Vec::value_type, NDIM> vector_key(const Vec& vec) {
if (vec.size() > NDIM) {
throw std::runtime_error(
fmt::format("ndim can not be larger than {}.", NDIM));
}
std::array<T, NDIM> result = {};
std::array<typename Vec::value_type, NDIM> result = {};
std::copy_n(vec.begin(), vec.size(), result.begin());
return result;
}
mlx/backend/cuda/device.cpp
+5 -2
View File
@@ -409,14 +409,17 @@ void CommandEncoder::commit() {
}
if (use_cuda_graphs() && node_count_ > 0) {
if (!from_nodes_.empty()) {
#if CUDART_VERSION >= 13000
CHECK_CUDA_ERROR(cudaGraphAddDependencies(
graph_,
from_nodes_.data(),
to_nodes_.data(),
#if CUDART_VERSION >= 13000
nullptr, // edgeData
#endif // CUDART_VERSION >= 13000
from_nodes_.size()));
#else
CHECK_CUDA_ERROR(cudaGraphAddDependencies(
graph_, from_nodes_.data(), to_nodes_.data(), from_nodes_.size()));
#endif
}
device_.make_current();
mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu
+5 -5
View File
@@ -160,7 +160,7 @@ class GemmGroupedEncoder
void encode(cu::CommandEncoder& encoder) {
encoder.add_kernel_node(
cutlass::Kernel<GemmKernel>,
{static_cast<uint>(this->params_.threadblock_count), 1, 1},
{static_cast<uint32_t>(this->params_.threadblock_count), 1, 1},
{GemmKernel::kThreadCount, 1, 1},
sizeof(typename GemmKernel::SharedStorage),
this->params_);
@@ -184,11 +184,11 @@ void grouped_gemm_v2(
dispatch_bool(a_transposed, [&](auto a_transposed_tag) {
dispatch_bool(b_transposed, [&](auto b_transposed_tag) {
using LayoutA = std::conditional_t<
a_transposed_tag,
a_transposed_tag.value,
cutlass::layout::ColumnMajor,
cutlass::layout::RowMajor>;
using LayoutB = std::conditional_t<
b_transposed_tag,
b_transposed_tag.value,
cutlass::layout::ColumnMajor,
cutlass::layout::RowMajor>;
using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped<
@@ -302,9 +302,9 @@ void cutlass_grouped_gemm_unaligned(
// Fill the pointers by computing offsets from indices.
constexpr int N_READS = 4;
size_t n_threads = cuda::ceil_div(indices.size(), N_READS);
int n_threads = cuda::ceil_div(indices.size(), N_READS);
n_threads = group_count < n_threads ? n_threads : group_count;
dim3 block_dims(std::min(n_threads, 1024ul));
dim3 block_dims(std::min(n_threads, 1024));
dim3 num_blocks(1);
encoder.set_input_array(indices);
mlx/backend/cuda/jit_module.cpp
+2 -3
View File
@@ -12,7 +12,6 @@
#include <fmt/format.h>
#include <nvrtc.h>
#include <unistd.h>
namespace mlx::core::cu {
@@ -321,7 +320,7 @@ void load_module(
const std::string& ptx,
const std::vector<std::pair<std::string, std::string>>& ptx_kernels,
CUmodule& module_,
std::unordered_map<std::string, std::tuple<CUfunction, bool, uint>>&
std::unordered_map<std::string, std::tuple<CUfunction, bool, uint32_t>>&
kernels) {
// Load module.
char jit_log[4089] = {};
@@ -383,7 +382,7 @@ JitModule::~JitModule() {
CHECK_CUDA_ERROR(cuModuleUnload(module_));
}
std::pair<CUfunction, uint> JitModule::get_kernel_and_dims(
std::pair<CUfunction, uint32_t> JitModule::get_kernel_and_dims(
const std::string& kernel_name,
std::function<void(CUfunction)> configure_kernel) {
auto it = kernels_.find(kernel_name);
mlx/backend/cuda/jit_module.h
+3 -2
View File
@@ -99,13 +99,14 @@ class JitModule {
CUfunction get_kernel(
const std::string& kernel_name,
std::function<void(CUfunction)> configure_kernel = nullptr);
std::pair<CUfunction, uint> get_kernel_and_dims(
std::pair<CUfunction, uint32_t> get_kernel_and_dims(
const std::string& kernel_name,
std::function<void(CUfunction)> configure_kernel = nullptr);
private:
CUmodule module_{nullptr};
std::unordered_map<std::string, std::tuple<CUfunction, bool, uint>> kernels_;
std::unordered_map<std::string, std::tuple<CUfunction, bool, uint32_t>>
kernels_;
};
std::unordered_map<std::string, JitModule>& get_jit_module_cache();
mlx/backend/cuda/kernel_utils.cu
+3 -3
View File
@@ -30,15 +30,15 @@ std::pair<dim3, dim3> get_grid_and_block(int dim0, int dim1, int dim2) {
return std::make_pair(dim3(gx, gy, gz), dim3(bx, by, bz));
}
std::tuple<dim3, uint> get_launch_args(
std::tuple<dim3, uint32_t> get_launch_args(
size_t size,
const Shape& shape,
const Strides& strides,
bool large,
int work_per_thread /* = 1 */,
uint max_block_dim /* = 1024 */) {
uint32_t max_block_dim /* = 1024 */) {
size_t nthreads = cuda::ceil_div(size, work_per_thread);
uint block_dim = max_block_dim < nthreads ? max_block_dim : nthreads;
uint32_t block_dim = max_block_dim < nthreads ? max_block_dim : nthreads;
dim3 num_blocks;
if (large) {
num_blocks = get_2d_grid_dims(shape, strides, work_per_thread);
mlx/backend/cuda/kernel_utils.cuh
+4 -4
View File
@@ -123,19 +123,19 @@ std::pair<dim3, dim3> get_grid_and_block(int dim0, int dim1, int dim2);
// Get the num_blocks and block_dims assuming each thread handles
// |work_per_thread| elements of |arr|.
std::tuple<dim3, uint> get_launch_args(
std::tuple<dim3, uint32_t> get_launch_args(
size_t size,
const Shape& shape,
const Strides& strides,
bool large,
int work_per_thread = 1,
uint max_block_dim = 1024);
uint32_t max_block_dim = 1024);
inline std::tuple<dim3, uint> get_launch_args(
inline std::tuple<dim3, uint32_t> get_launch_args(
const array& arr,
bool large,
int work_per_thread = 1,
uint max_block_dim = 1024) {
uint32_t max_block_dim = 1024) {
return get_launch_args(
arr.size(),
arr.shape(),
mlx/backend/cuda/quantized/affine_quantize.cu
+1 -1
View File
@@ -208,7 +208,7 @@ __global__ void affine_dequantize(
bias;
out[3] = static_cast<T>((w[2] >> 2) & 0x3f) * scale + bias;
} else {
uint val = w[offset];
uint32_t val = w[offset];
#pragma clang loop unroll(full)
for (int i = 0; i < pack_factor; i++) {
uint8_t d;
mlx/backend/cuda/quantized/fp_quantize.cu
+1 -1
View File
@@ -247,7 +247,7 @@ fp_dequantize(const uint8_t* w, const uint8_t* scales, T* out, size_t size) {
out += oindex;
uint val = w[offset];
uint32_t val = w[offset];
#pragma clang loop unroll(full)
for (int i = 0; i < pack_factor; i++) {
uint8_t d;
mlx/backend/cuda/random.cu
+6 -6
View File
@@ -51,9 +51,9 @@ __global__ void rbitsc(
bool odd,
uint32_t bytes_per_key) {
auto grid = cg::this_grid();
uint thread_index = grid.thread_rank();
uint index_x = thread_index % grid_dims.x;
uint index_y = thread_index / grid_dims.x;
uint32_t thread_index = grid.thread_rank();
uint32_t index_x = thread_index % grid_dims.x;
uint32_t index_y = thread_index / grid_dims.x;
if (index_x >= grid_dims.x || index_y >= grid_dims.y) {
return;
}
@@ -94,9 +94,9 @@ __global__ void rbits(
const __grid_constant__ Shape key_shape,
const __grid_constant__ Strides key_strides) {
auto grid = cg::this_grid();
uint thread_index = grid.thread_rank();
uint index_x = thread_index % grid_dims.x;
uint index_y = thread_index / grid_dims.x;
uint32_t thread_index = grid.thread_rank();
uint32_t index_x = thread_index % grid_dims.x;
uint32_t index_y = thread_index / grid_dims.x;
if (index_x >= grid_dims.x || index_y >= grid_dims.y) {
return;
}
mlx/backend/cuda/reduce/all_reduce.cu
+2 -2
View File
@@ -68,8 +68,8 @@ void all_reduce(
out.set_data(cu::malloc_async(out.nbytes(), encoder));
auto get_args = [](size_t size, int N) {
int threads = std::min(512UL, (size + N - 1) / N);
auto get_args = [](int size, int N) {
int threads = std::min(512, (size + N - 1) / N);
threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
int reductions_per_step = threads * N;
size_t steps_needed =
mlx/backend/cuda/rope.cu
+1 -1
View File
@@ -30,7 +30,7 @@ __device__ void rope_single_impl(
float sintheta = sin(theta);
// Compute the input and output indices
uint index_1, index_2;
uint32_t index_1, index_2;
if (traditional) {
index_1 = 2 * pos.x + pos.y * stride;
index_2 = index_1 + 1;
mlx/backend/cuda/scaled_dot_product_attention.cu
+3
View File
@@ -1,5 +1,8 @@
// Copyright © 2025 Apple Inc.
// Required for using M_LOG2E in MSVC.
#define _USE_MATH_DEFINES
#include "mlx/backend/cuda/device.h"
#include "mlx/backend/cuda/device/config.h"
#include "mlx/backend/cuda/device/utils.cuh"
mlx/backend/cuda/scan.cu
+9 -9
View File
@@ -227,21 +227,21 @@ __global__ void strided_scan(
// Compute offsets.
int64_t offset = (grid.block_rank() / stride_blocks) * axis_size * stride;
int64_t global_index_x = (grid.block_rank() % stride_blocks) * BN;
uint read_offset_y = (block.thread_rank() * N_READS) / BN;
uint read_offset_x = (block.thread_rank() * N_READS) % BN;
uint scan_offset_y = warp.thread_rank();
uint scan_offset_x = warp.meta_group_rank() * n_scans;
uint32_t read_offset_y = (block.thread_rank() * N_READS) / BN;
uint32_t read_offset_x = (block.thread_rank() * N_READS) % BN;
uint32_t scan_offset_y = warp.thread_rank();
uint32_t scan_offset_x = warp.meta_group_rank() * n_scans;
uint stride_limit = stride - global_index_x;
uint32_t stride_limit = stride - global_index_x;
in += offset + global_index_x + read_offset_x;
out += offset + global_index_x + read_offset_x;
U* read_into = read_buffer + read_offset_y * BN_pad + read_offset_x;
U* read_from = read_buffer + scan_offset_y * BN_pad + scan_offset_x;
for (uint j = 0; j < axis_size; j += BM) {
for (uint32_t j = 0; j < axis_size; j += BM) {
// Calculate the indices for the current thread.
uint index_y = j + read_offset_y;
uint check_index_y = index_y;
uint32_t index_y = j + read_offset_y;
uint32_t check_index_y = index_y;
if (reverse) {
index_y = axis_size - 1 - index_y;
}
@@ -395,7 +395,7 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
dispatch_scan_ops(reduce_type_, [&](auto scan_op_tag) {
using Op = MLX_GET_TYPE(scan_op_tag);
if constexpr (supports_scan_op<Op, T>) {
if constexpr (supports_scan_op<Op, T>()) {
using U = typename cu::ScanResult<Op, T>::type;
dispatch_bool(inclusive_, [&](auto inclusive) {
dispatch_bool(reverse_, [&](auto reverse) {
mlx/backend/cuda/utils.h
+1 -1
View File
@@ -11,7 +11,7 @@
namespace mlx::core {
template <typename T>
inline uint max_occupancy_block_dim(T kernel) {
inline uint32_t max_occupancy_block_dim(T kernel) {
int _, block_dim;
if constexpr (std::is_same_v<T, CUfunction>) {
CHECK_CUDA_ERROR(
mlx/primitives.h
+1 -1
View File
@@ -46,7 +46,7 @@
namespace mlx::core {
// Abstract base class
class Primitive {
class MLX_API Primitive {
public:
explicit Primitive(Stream stream) : stream_(stream) {}