add kernel with torch

This commit is contained in:
Pritimay Sarkar
2023-12-10 16:49:32 +05:30
parent a62238b30d
commit c0b9a88a53
2 changed files with 44 additions and 0 deletions

30
cuda/cudaprog.cu Normal file
View File

@@ -0,0 +1,30 @@
#include <ATen/ATen.h>
__global__ void my_cuda_kernel(float* input, float* output, int size) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < size) {
output[idx] = 2 * input[idx];
}
}
at::Tensor my_cuda_function(at::Tensor input) {
int size = input.numel();
// Allocate output tensor on the GPU
at::Tensor output = at::empty_like(input);
// Launch the CUDA kernel
dim3 blockDim(256);
dim3 gridDim((size + blockDim.x - 1) / blockDim.x);
my_cuda_kernel<<<gridDim, blockDim>>>(
input.data<float>(),
output.data<float>(),
size
);
// Wait for the kernel to finish
cudaDeviceSynchronize();
return output;
}