aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

batched matmul cuda

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-16 19:53:39 -03:00
parent 5da5a7ad0d
commit b3c0543114
1 changed files with 1 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

2
norch/csrc/cuda.cu
View file

@ -402,7 +402,7 @@ __host__ void batched_matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float
int cols2 = tensor2->shape[2];
dim3 threadsPerBlock(16, 16);
dim3 number_of_blocks((cols2 + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows1 + threadsPerBlock.y - 1) / threadsPerBlock.y);
dim3 number_of_blocks((cols2 + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows1 + threadsPerBlock.y - 1) / threadsPerBlock.y, batch_size);
batched_matmul_tensor_cuda_kernel<<<number_of_blocks, threadsPerBlock>>>(tensor1->data, tensor2->data, result_data, batch_size, rows1, cols1, cols2);