zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-02 01:08:10 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

[CUDA] Matmul utils initial commit (#2441)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-08-01 14:22:25 -07:00
committed by GitHub
parent 86258f292f
commit be9bc96da4
32 changed files with 856 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -129,7 +129,7 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
auto kernel = mod.get_kernel(kernel_name);
auto [num_blocks, block_dims] = get_launch_args(out, large);
encoder.add_kernel_node(kernel, num_blocks, block_dims, args.args());
encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
}
void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -230,7 +230,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(upd, large);
encoder.add_kernel_node(kernel, num_blocks, block_dims, args.args());
encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
}
void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -318,7 +318,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(idx, large);
encoder.add_kernel_node(kernel, num_blocks, block_dims, args.args());
encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
}
void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -422,7 +422,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(idx, large);
encoder.add_kernel_node(kernel, num_blocks, block_dims, args.args());
encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
}
} // namespace mlx::core