Puzzle 3: 가드

개요

벡터 a의 각 위치에 10을 더해 output에 저장하는 커널을 구현해 보세요.

참고: 스레드 수가 데이터 개수보다 많아서, 일부 스레드는 처리할 데이터가 없습니다. 이런 스레드가 범위를 벗어난 메모리에 접근하지 않도록 방지해야 합니다.

핵심 개념

이 퍼즐에서 다루는 내용:

• 스레드 수와 데이터 크기 불일치 처리
• 범위를 벗어난 메모리 접근 방지
• GPU 커널에서 조건부 실행 사용
• 안전한 메모리 접근 패턴

수학적 표현

각 스레드 \(i\)에 대해: \[\Large \text{if}\ i < \text{size}: output[i] = a[i] + 10\]

메모리 안전 패턴

Thread 0 (i=0):  if 0 < size:  output[0] = a[0] + 10  ✓ Valid
Thread 1 (i=1):  if 1 < size:  output[1] = a[1] + 10  ✓ Valid
Thread 2 (i=2):  if 2 < size:  output[2] = a[2] + 10  ✓ Valid
Thread 3 (i=3):  if 3 < size:  output[3] = a[3] + 10  ✓ Valid
Thread 4 (i=4):  if 4 < size:  ❌ Skip (out of bounds)
Thread 5 (i=5):  if 5 < size:  ❌ Skip (out of bounds)

💡 참고: 다음 상황에서 경계(boundary) 검사는 점점 복잡해집니다:

• 다차원 배열
• 다양한 배열 형태
• 복잡한 접근 패턴

완성할 코드

comptime SIZE = 4
comptime BLOCKS_PER_GRID = 1
comptime THREADS_PER_BLOCK = 8
comptime dtype = DType.float32


fn add_10_guard(
    output: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    a: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    size: UInt,
):
    i = thread_idx.x
    # FILL ME IN (roughly 2 lines)

전체 코드 보기: problems/p03/p03.mojo

코드 실행

솔루션을 테스트하려면 터미널에서 다음 명령어를 실행하세요:

pixi run p03

pixi run -e amd p03

pixi run -e apple p03

uv run poe p03

퍼즐을 아직 풀지 않았다면 출력이 다음과 같이 나타납니다:

out: HostBuffer([0.0, 0.0, 0.0, 0.0])
expected: HostBuffer([10.0, 11.0, 12.0, 13.0])

솔루션

fn add_10_guard(
    output: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    a: UnsafePointer[Scalar[dtype], MutAnyOrigin],
    size: UInt,
):
    i = thread_idx.x
    if i < size:
        output[i] = a[i] + 10.0

앞으로 다룰 내용

간단한 경계 검사는 여기서 잘 작동하지만, 다음 상황을 생각해 보세요:

• 2D/3D 배열의 경계는 어떻게 처리할까?
• 다양한 형태를 효율적으로 처리하려면?
• 패딩이나 가장자리 처리가 필요하다면?

복잡도가 증가하는 예시:

# 현재: 1D 경계 검사
if i < size: ...

# 곧 다룰 내용: 2D 경계 검사
if i < height and j < width: ...

# 이후: 패딩이 있는 3D
if i < height and j < width and k < depth and
   i >= padding and j >= padding: ...

이런 경계 처리 패턴은 Puzzle 4의 LayoutTensor 알아보기에서 배우면 훨씬 깔끔해집니다. LayoutTensor는 형태 관리 기능을 기본으로 제공합니다.