zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-07 04:06:42 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Remove the kernel arg from get_launch_args (#2437)

Browse Source
This commit is contained in:
Cheng 2025-07-30 11:43:02 +09:00 committed by GitHub
parent 3adba92ebe
commit 254476718b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 83 additions and 125 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
mlx/backend/cuda/binary.cu
View File

@ -211,12 +211,15 @@ void binary_op_gpu_inplace(
int ndim = shape.size();
if (ndim <= 3) {
dispatch_1_2_3(ndim, [&](auto dims_constant) {
auto kernel = cu::
binary_g_nd<Op, InType, OutType, IdxT, dims_constant()>;
auto [num_blocks, block_dims] =
get_launch_args(kernel, out, large());
get_launch_args(out, large());
encoder.add_kernel_node(
kernel,
cu::binary_g_nd<
Op,
InType,
OutType,
IdxT,
dims_constant()>,
num_blocks,
block_dims,
a.data<InType>(),
@ -228,11 +231,9 @@ void binary_op_gpu_inplace(
const_param<dims_constant()>(b_strides));
});
} else {
auto kernel = cu::binary_g<Op, InType, OutType, IdxT>;
auto [num_blocks, block_dims] =
get_launch_args(kernel, out, large());
auto [num_blocks, block_dims] = get_launch_args(out, large());
encoder.add_kernel_node(
kernel,
cu::binary_g<Op, InType, OutType, IdxT>,
num_blocks,
block_dims,
a.data<InType>(),
@ -258,12 +259,7 @@ void binary_op_gpu_inplace(
kernel = cu::binary_vv<Op, InType, OutType, IdxT, N_READS>;
}
auto [num_blocks, block_dims] = get_launch_args(
kernel,
out.data_size(),
out.shape(),
out.strides(),
large(),
N_READS);
out.data_size(), out.shape(), out.strides(), large(), N_READS);
encoder.add_kernel_node(
kernel,
num_blocks,

21
mlx/backend/cuda/binary_two.cu
View File

@ -227,16 +227,15 @@ void binary_two_op_gpu_inplace(
int ndim = shape.size();
if (ndim <= 3) {
dispatch_1_2_3(ndim, [&](auto dims_constant) {
auto kernel = cu::binary_two_g_nd<
Op,
InType,
OutType,
IdxT,
dims_constant()>;
auto [num_blocks, block_dims] =
get_launch_args(kernel, out_a, large());
get_launch_args(out_a, large());
encoder.add_kernel_node(
kernel,
cu::binary_two_g_nd<
Op,
InType,
OutType,
IdxT,
dims_constant()>,
num_blocks,
block_dims,
a.data<InType>(),
@ -249,11 +248,10 @@ void binary_two_op_gpu_inplace(
const_param<dims_constant()>(b_strides));
});
} else {
auto kernel = cu::binary_two_g<Op, InType, OutType, IdxT>;
auto [num_blocks, block_dims] =
get_launch_args(kernel, out_a, large());
get_launch_args(out_a, large());
encoder.add_kernel_node(
kernel,
cu::binary_two_g<Op, InType, OutType, IdxT>,
num_blocks,
block_dims,
a.data<InType>(),
@ -280,7 +278,6 @@ void binary_two_op_gpu_inplace(
kernel = cu::binary_two_vv<Op, InType, OutType, IdxT, N_READS>;
}
auto [num_blocks, block_dims] = get_launch_args(
kernel,
out_a.data_size(),
out_a.shape(),
out_a.strides(),

2
mlx/backend/cuda/compiled.cpp
View File

@ -294,7 +294,7 @@ void Compiled::eval_gpu(
auto kernel = mod.get_kernel(kernel_name);
auto [num_blocks, block_dims] =
get_launch_args(kernel, outputs[0], large, work_per_thread);
get_launch_args(outputs[0], large, work_per_thread);
encoder.add_kernel_node(kernel, num_blocks, block_dims, args.args());
}

7
mlx/backend/cuda/copy/copy_contiguous.cu
View File

@ -71,12 +71,7 @@ void copy_contiguous(
kernel = cu::copy_v<InType, OutType, IdxT, N_READS>;
}
auto [num_blocks, block_dims] = get_launch_args(
kernel,
out.data_size(),
out.shape(),
out.strides(),
large(),
N_READS);
out.data_size(), out.shape(), out.strides(), large(), N_READS);
encoder.add_kernel_node(
kernel,
num_blocks,

15
mlx/backend/cuda/copy/copy_general.cu
View File

@ -71,12 +71,10 @@ void copy_general(
data_size *= s;
if (ndim <= 3) {
dispatch_1_2_3(ndim, [&](auto ndim_constant) {
auto kernel =
cu::copy_gg_nd<InType, OutType, IdxT, ndim_constant()>;
auto [num_blocks, block_dims] = get_launch_args(
kernel, data_size, shape, out.strides(), large());
auto [num_blocks, block_dims] =
get_launch_args(data_size, shape, out.strides(), large());
encoder.add_kernel_node(
kernel,
cu::copy_gg_nd<InType, OutType, IdxT, ndim_constant()>,
num_blocks,
block_dims,
in_ptr,
@ -87,11 +85,10 @@ void copy_general(
const_param<ndim_constant()>(strides_out));
});
} else { // ndim >= 4
auto kernel = cu::copy_gg<InType, OutType, IdxT>;
auto [num_blocks, block_dims] = get_launch_args(
kernel, data_size, shape, out.strides(), large());
auto [num_blocks, block_dims] =
get_launch_args(data_size, shape, out.strides(), large());
encoder.add_kernel_node(
kernel,
cu::copy_gg<InType, OutType, IdxT>,
num_blocks,
block_dims,
in_ptr,

17
mlx/backend/cuda/copy/copy_general_dynamic.cu
View File

@ -74,12 +74,13 @@ void copy_general_dynamic(
int ndim = shape.size();
if (ndim <= 3) {
dispatch_1_2_3(ndim, [&](auto dims_constant) {
auto kernel = cu::
copy_gg_dynamic_nd<InType, OutType, IdxT, dims_constant()>;
auto [num_blocks, block_dims] =
get_launch_args(kernel, out, large());
auto [num_blocks, block_dims] = get_launch_args(out, large());
encoder.add_kernel_node(
kernel,
cu::copy_gg_dynamic_nd<
InType,
OutType,
IdxT,
dims_constant()>,
num_blocks,
block_dims,
in_ptr,
@ -92,11 +93,9 @@ void copy_general_dynamic(
dynamic_offset_out.data<int64_t>());
});
} else { // ndim >= 4
auto kernel = cu::copy_gg_dynamic<InType, OutType, IdxT>;
auto [num_blocks, block_dims] =
get_launch_args(kernel, out, large());
auto [num_blocks, block_dims] = get_launch_args(out, large());
encoder.add_kernel_node(
kernel,
cu::copy_gg_dynamic<InType, OutType, IdxT>,
num_blocks,
block_dims,
in_ptr,

13
mlx/backend/cuda/copy/copy_general_input.cu
View File

@ -63,12 +63,9 @@ void copy_general_input(
int ndim = shape.size();
if (ndim <= 3) {
dispatch_1_2_3(ndim, [&](auto dims_constant) {
auto kernel =
cu::copy_g_nd<InType, OutType, IdxT, dims_constant()>;
auto [num_blocks, block_dims] =
get_launch_args(kernel, out, large());
auto [num_blocks, block_dims] = get_launch_args(out, large());
encoder.add_kernel_node(
kernel,
cu::copy_g_nd<InType, OutType, IdxT, dims_constant()>,
num_blocks,
block_dims,
in_ptr,
@ -78,11 +75,9 @@ void copy_general_input(
const_param<dims_constant()>(strides_in));
});
} else { // ndim >= 4
auto kernel = cu::copy_g<InType, OutType, IdxT>;
auto [num_blocks, block_dims] =
get_launch_args(kernel, out, large());
auto [num_blocks, block_dims] = get_launch_args(out, large());
encoder.add_kernel_node(
kernel,
cu::copy_g<InType, OutType, IdxT>,
num_blocks,
block_dims,
in_ptr,

8
mlx/backend/cuda/indexing.cpp
View File

@ -128,7 +128,7 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
encoder.set_output_array(out);
auto kernel = mod.get_kernel(kernel_name);
auto [num_blocks, block_dims] = get_launch_args(kernel, out, large);
auto [num_blocks, block_dims] = get_launch_args(out, large);
encoder.add_kernel_node(kernel, num_blocks, block_dims, args.args());
}
@ -229,7 +229,7 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
}
encoder.set_output_array(out);
auto kernel = mod.get_kernel(kernel_name);
auto [num_blocks, block_dims] = get_launch_args(kernel, upd, large);
auto [num_blocks, block_dims] = get_launch_args(upd, large);
encoder.add_kernel_node(kernel, num_blocks, block_dims, args.args());
}
@ -317,7 +317,7 @@ void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
}
encoder.set_output_array(out);
auto kernel = mod.get_kernel(kernel_name);
auto [num_blocks, block_dims] = get_launch_args(kernel, idx, large);
auto [num_blocks, block_dims] = get_launch_args(idx, large);
encoder.add_kernel_node(kernel, num_blocks, block_dims, args.args());
}
@ -421,7 +421,7 @@ void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
}
encoder.set_output_array(out);
auto kernel = mod.get_kernel(kernel_name);
auto [num_blocks, block_dims] = get_launch_args(kernel, idx, large);
auto [num_blocks, block_dims] = get_launch_args(idx, large);
encoder.add_kernel_node(kernel, num_blocks, block_dims, args.args());
}

21
mlx/backend/cuda/kernel_utils.cu
View File

@ -30,4 +30,25 @@ 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(
size_t size,
const Shape& shape,
const Strides& strides,
bool large,
int work_per_thread) {
size_t nthreads = cuda::ceil_div(size, work_per_thread);
uint block_dim = 1024;
if (block_dim > nthreads) {
block_dim = nthreads;
}
dim3 num_blocks;
if (large) {
num_blocks = get_2d_grid_dims(shape, strides, work_per_thread);
num_blocks.x = cuda::ceil_div(num_blocks.x, block_dim);
} else {
num_blocks.x = cuda::ceil_div(nthreads, block_dim);
}
return std::make_tuple(num_blocks, block_dim);
}
} // namespace mlx::core

30
mlx/backend/cuda/kernel_utils.cuh
View File

@ -122,37 +122,17 @@ std::pair<dim3, dim3> get_grid_and_block(int dim0, int dim1, int dim2);
// Get the num_blocks and block_dims that maximize occupancy for |kernel|,
// assuming each thread handles |work_per_thread| elements of |arr|.
template <typename T>
inline std::tuple<dim3, uint> get_launch_args(
T kernel,
std::tuple<dim3, uint> get_launch_args(
size_t size,
const Shape& shape,
const Strides& strides,
bool large,
int work_per_thread = 1) {
size_t nthreads = cuda::ceil_div(size, work_per_thread);
uint block_dim = 1024;
if (block_dim > nthreads) {
block_dim = nthreads;
}
dim3 num_blocks;
if (large) {
num_blocks = get_2d_grid_dims(shape, strides, work_per_thread);
num_blocks.x = cuda::ceil_div(num_blocks.x, block_dim);
} else {
num_blocks.x = cuda::ceil_div(nthreads, block_dim);
}
return std::make_tuple(num_blocks, block_dim);
}
int work_per_thread = 1);
template <typename T>
inline std::tuple<dim3, uint> get_launch_args(
T kernel,
const array& arr,
bool large,
int work_per_thread = 1) {
inline std::tuple<dim3, uint>
get_launch_args(const array& arr, bool large, int work_per_thread = 1) {
return get_launch_args(
kernel, arr.size(), arr.shape(), arr.strides(), large, work_per_thread);
arr.size(), arr.shape(), arr.strides(), large, work_per_thread);
}
} // namespace mlx::core

12
mlx/backend/cuda/quantized.cu
View File

@ -350,12 +350,10 @@ void fast::AffineQuantize::eval_gpu(
dispatch_bits(bits_, [&](auto bits) {
using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
if (dequantize_) {
auto kernel =
cu::affine_dequantize<DataType, group_size.value, bits.value>;
auto [num_blocks, block_dims] =
get_launch_args(kernel, size, grid_shape, w.strides(), large);
get_launch_args(size, grid_shape, w.strides(), large);
enc.add_kernel_node(
kernel,
cu::affine_dequantize<DataType, group_size.value, bits.value>,
num_blocks,
block_dims,
w.data<uint8_t>(),
@ -364,12 +362,10 @@ void fast::AffineQuantize::eval_gpu(
out.data<DataType>(),
out.size());
} else {
auto kernel =
cu::affine_quantize<DataType, group_size.value, bits.value>;
auto [num_blocks, block_dims] =
get_launch_args(kernel, size, grid_shape, w.strides(), large);
get_launch_args(size, grid_shape, w.strides(), large);
enc.add_kernel_node(
kernel,
cu::affine_quantize<DataType, group_size.value, bits.value>,
num_blocks,
block_dims,
w.data<DataType>(),

23
mlx/backend/cuda/ternary.cu
View File

@ -125,12 +125,9 @@ void ternary_op_gpu_inplace(
int ndim = shape.size();
if (ndim <= 3) {
dispatch_1_2_3(ndim, [&](auto dims_constant) {
auto kernel =
cu::ternary_g_nd<Op, DType, IdxT, dims_constant()>;
auto [num_blocks, block_dims] =
get_launch_args(kernel, out, large());
auto [num_blocks, block_dims] = get_launch_args(out, large());
encoder.add_kernel_node(
kernel,
cu::ternary_g_nd<Op, DType, IdxT, dims_constant()>,
num_blocks,
block_dims,
a.data<bool>(),
@ -144,11 +141,9 @@ void ternary_op_gpu_inplace(
const_param<dims_constant()>(c_strides));
});
} else {
auto kernel = cu::ternary_g<Op, DType, IdxT>;
auto [num_blocks, block_dims] =
get_launch_args(kernel, out, large());
auto [num_blocks, block_dims] = get_launch_args(out, large());
encoder.add_kernel_node(
kernel,
cu::ternary_g<Op, DType, IdxT>,
num_blocks,
block_dims,
a.data<bool>(),
@ -167,16 +162,10 @@ void ternary_op_gpu_inplace(
dispatch_bool(out.data_size() > UINT32_MAX, [&](auto large) {
using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
constexpr int N_READS = 16 / sizeof(DType);
auto kernel = cu::ternary_v<Op, DType, IdxT, N_READS>;
auto [num_blocks, block_dims] = get_launch_args(
kernel,
out.data_size(),
out.shape(),
out.strides(),
large(),
N_READS);
out.data_size(), out.shape(), out.strides(), large(), N_READS);
encoder.add_kernel_node(
kernel,
cu::ternary_v<Op, DType, IdxT, N_READS>,
num_blocks,
block_dims,
a.data<bool>(),

15
mlx/backend/cuda/unary.cu
View File

@ -129,16 +129,10 @@ void unary_op_gpu_inplace(
using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
// TODO: Choose optimized value based on type size.
constexpr int N_READS = 4;
auto kernel = cu::unary_v<Op, InType, OutType, IdxT, N_READS>;
auto [num_blocks, block_dims] = get_launch_args(
kernel,
out.data_size(),
out.shape(),
out.strides(),
large,
N_READS);
out.data_size(), out.shape(), out.strides(), large, N_READS);
encoder.add_kernel_node(
kernel,
cu::unary_v<Op, InType, OutType, IdxT, N_READS>,
num_blocks,
block_dims,
in.data<InType>(),
@ -147,10 +141,9 @@ void unary_op_gpu_inplace(
} else {
using IdxT = std::conditional_t<large(), int64_t, int32_t>;
auto [shape, strides] = collapse_contiguous_dims(in);
auto kernel = cu::unary_g<Op, InType, OutType, IdxT>;
auto [num_blocks, block_dims] = get_launch_args(kernel, out, large);
auto [num_blocks, block_dims] = get_launch_args(out, large);
encoder.add_kernel_node(
kernel,
cu::unary_g<Op, InType, OutType, IdxT>,
num_blocks,
block_dims,
in.data<InType>(),