Easy
Implement a program that performs element-wise addition of two vectors containing 32-bit floating point numbers on a GPU. The program should take two input vectors of equal length and produce a single output vector containing their sum.
Implementation Requirements
External libraries are not permitted
The solve function signature must remain unchanged
The final result must be stored in vector C
Example 1:
Input: A = [1.0, 2.0, 3.0, 4.0]
B = [5.0, 6.0, 7.0, 8.0]
Output: C = [6.0, 8.0, 10.0, 12.0]
Example 2:
Input: A = [1.5, 1.5, 1.5]
B = [2.3, 2.3, 2.3]
Output: C = [3.8, 3.8, 3.8]
Constraints
Input vectors A and B have identical lengths
1 ≤ N ≤ 100,000,000
import torch
# 메모리 효율적 버전
# A, B, C are tensors on the GPU
def solve(
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
N: int
):
"""
1. pytorch가 적절한 cuda 커널 선택
- 텐서크기, 데이터타입 등 고려
- elementwisr_add_kernel 같은 커널 선택
2. gpu 메모리 할당
- 결과를저장할 메모리 공간 확보
3. 커널 호출 파라미터 계산
- 블록수 = (N + 255) / 256
- 스레드 수 = 256
4. CUDA 커널 실행
- 각 스레드가 하나 이상의 요소 처리
- 모든 연산이 병렬로 수행
5. 동기화 및 결과 반환
- gpu 연산 완료 대기
- 결과 텐서 반환
"""
# element-wise addition of two vectors 32-bit floating point
# equal length
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 큰 벡터의 경우 청크 단위로 처리
chunk_size = 10000000
for i in range(0, N, chunk_size):
end_idx = min(i + chunk_size, N)
# 입력 데이터를 pytorch tensor로 변환하고 gpu로 이동
chunk_A = torch.tensor(A[i:end_idx], dtype=torch.float32, device=device)
chunk_B = torch.tensor(B[i:end_idx], dtype=torch.float32, device=device)
# gpu에서 요소별 덧셈 수행
# pytorch는 내부적으로 최적화된 cuda 커널을 호출
chunk_C = chunk_A + chunk_B
# 결과물을 cpu로 복사하고 python리스트로 변환
result_chunk = chunk_C.cpu().numpy()
for j, val in enumerate(result_chunk):
C[i + j] = float(val)
# Gpu 메모리 정리
del chunk_A, chunk_B, chunk_C
torch.cuda.empty_cache()
```
```
#include <cuda_runtime.h>
#include <stdio.h>
// 커널함수
__global__ void vector_add_optimized(const float* A, const float* B, float* C, int N) {
// Grid-stride loop: 각 스레드가 여러 요소 처리
// 각 스레드가 처리할 전역 인덱스 계산
int idx = blockIdx.x * blockDim.X + threadIdx.x;
// blockIdx.x : 현재 블록의 번호
// blockDim.x : 블록당 스레드 수 (여기서는 256)
// threadIdx.x : 블록 내에서 현재 스레드 번호
// 블록 0: 스레드 0~255가 벡터 인덱스 0~255 처리
// 블록 1: 스레드 0~255가 벡터 인덱스 256~511 처리
// 블록 2: 스레드 0~255가 벡터 인덱스 512~767 처리
int stride = blockDim.x * gridDim.x;
// // 경계 검사: 인덱스가 벡터 크기를 초과하지 않도록
// if (idx < N){
// // 각 스레드가 하나의 요소를 처리
// C[idx] = A[idx] + B[idx];
// }
for (int i = idx; i < N; i += stride){
C[i] = A[i] + B[i];
}
}
// -> 메모리 접근 최적화 버전 (float 4 사용)
__global__ void vector_add_float4(const float* A, const float* B, float* C, int N) {
int idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
if (idx + 3 < N) {
// 4개 요소를 한번에 로드/저장 (128 비트 단위)
float4 a = *reinterpret_cast<const float4*>(&A[idx]);
float4 b = *reinterpret_cast<const float4*>(&B[idx]);
float4 c;
c.x = a.x + b.x;
c.y = a.y + b.y;
c.z = a.z + b.z;
c.w = a.w + b.w;
*reinterpret_cast<float4*>(&C[idx]) = c;
}
}
//A, B, C are device pointers
//(i.e. pointers to memory on the GPU)
extern "C" void solve(const float* A, const float* B, float* C, int N) {
int threadsPerBlock = 256;
int blockPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;
vector_add<<<blocksPerGrid, threadsPerBlock>>>(A, B, C, N);
cudaDeviceSynchronize();
}
int main() {
int N = 10000000;
size_t size = N * sizeof(float);
// host memory 할당
float *h_A = (float*)malloc(size);
float *h_B = (float*)malloc(size);
float *h_C = (float*)malloc(size);
// 초기화
for (int i = 0; i < N; i ++) {
h_A[i] = 1.0f;
h_B[i] = 2.0f;
}
// 디바이스 메모리 할당
float *d_A, *d_B, *d_C;
cudaMalloc(&d_A, size);
cudaMalloc(&d_B, size);
cudaMalloc(&d_C, size);
// 호스트 -> 디바이스 복사
cudaMemcpy(d_A, h_A, size, cudaMemcpyHostToDevice);
cudaMemcpy(d_B, h_B, size, cudaMemcpyHostToDevice);
// solve 함수 호출
solve(d_A, d_B, d_C, N);
// 결과를 호스트로 복사
cudaMemcpy(h_C, d_C, size, cudaMemcpyDeviceToHost);
//결과 확인
printf("first 5 results: ");
for (int i = 0; i < 5; i ++) {
printf("%.1f ", h_C[i]);
}
printf("\n")
// memory 해제
free(h_A); free(h_B); free(h_C);
cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
return 0;
}
```