向量加法

2026-09-26 22:21:54 +08:00
commit f57047a49b
3 changed files with 74 additions and 0 deletions
@@ -0,0 +1,13 @@
+#include <stdint.h>
+#include <cuda_fp16.h>
+__global__ void vector_add(__half* A, const __half* B, int N) {
+    int idx=blockDim.x*blockIdx.x+threadIdx.x;
+    if(idx<N) A[idx]=__hadd(A[idx],B[idx]);
+}
+extern "C" void run_kernel(__half* A, const __half* B, int64_t numel){
+    int threadsPerBlock = 256;
+    int blocksPerGrid = (numel + threadsPerBlock - 1) / threadsPerBlock;
+
+    vector_add<<<blocksPerGrid, threadsPerBlock>>>(A, B, numel);
+    cudaDeviceSynchronize();
+}