向量加法

This commit is contained in:
2026-09-26 22:21:54 +08:00
commit f57047a49b
3 changed files with 74 additions and 0 deletions
+13
View File
@@ -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();
}
+27
View File
@@ -0,0 +1,27 @@
#include <stdint.h>
#include <cuda_fp16.h>
__global__ void vector_add(__half* A, const __half* B, int numel) {
int idx=blockDim.x*blockIdx.x+threadIdx.x;
__half2* A2 = reinterpret_cast< __half2* >(A);
const __half2* B2 = reinterpret_cast<const __half2* >(B);
int base = blockIdx.x * blockDim.x * 4 + threadIdx.x;
int i0 = base;
int i1 = base + blockDim.x;
int i2 = base + blockDim.x * 2;
int i3 = base + blockDim.x * 3;
int n2 = numel/2;
if(i0 < n2) A2[i0] = __hadd2(A2[i0],B2[i0]);
if(i1 < n2) A2[i1] = __hadd2(A2[i1],B2[i1]);
if(i2 < n2) A2[i2] = __hadd2(A2[i2],B2[i2]);
if(i3 < n2) A2[i3] = __hadd2(A2[i3],B2[i3]);
if(idx == 0 && (numel&1)){
A[numel-1]= __hadd(A[numel-1],B[numel-1]);
}
}
extern "C" void run_kernel(__half* A, const __half* B, int64_t numel){
int threads = 1024;
int numel2=numel/2;
int works=(numel2+3)/4;
int blocks = (works + threads-1) / threads;
vector_add<<<blocks,threads>>>(A,B,numel);
}
+34
View File
@@ -0,0 +1,34 @@
#include <stdint.h>
#include <cuda_fp16.h>
struct __align__(16) half4 {
__half2 a, b, c, d;
};
__global__ void __launch_bounds__(256, 2) fp16_add_kernel(__half* __restrict__ A,
const __half* __restrict__ B,
int64_t numel) {
int64_t idx = ((int64_t)blockIdx.x * blockDim.x + threadIdx.x) << 3;
if (idx + 7 < numel) {
half4 a4 = *reinterpret_cast<half4*>(A + idx);
half4 b4 = *reinterpret_cast<const half4*>(B + idx);
a4.a = __hadd2(a4.a, b4.a);
a4.b = __hadd2(a4.b, b4.b);
a4.c = __hadd2(a4.c, b4.c);
a4.d = __hadd2(a4.d, b4.d);
*reinterpret_cast<half4*>(A + idx) = a4;
} else if (idx < numel) {
#pragma unroll
for (int i = 0; i < 8 && idx + i < numel; ++i)
A[idx + i] = __hadd(A[idx + i], B[idx + i]);
}
}
extern "C" void run_kernel(__half* A, const __half* B, int64_t numel) {
const int threads = 256;
int64_t blocks = (numel + ((int64_t)threads << 3) - 1) >> 11;
fp16_add_kernel<<<blocks, threads>>>(A, B, numel);
}