add kernel with torch

cuda/cudaprog.cu

@@ -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;
+}

train/run_cuda.py

@@ -0,0 +1,14 @@
+import torch
+from torch.utils.cpp_extension import load
+
+# Load the compiled CUDA extension
+my_cuda_extension = load(name='my_cuda_kernel', sources=['my_cuda_kernel.cu'])
+
+# Define a PyTorch tensor on the GPU
+input_tensor = torch.cuda.FloatTensor([1, 2, 3, 4])
+
+# Call the CUDA function
+output_tensor = my_cuda_extension.my_cuda_function(input_tensor)
+
+# Print the result
+print(output_tensor)