zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 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
13 changed files with 83 additions and 125 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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,