Puzzle 18: 소프트맥스 Op

fn softmax_gpu_kernel[
    layout: Layout,
    input_size: Int,
    dtype: DType = DType.float32,
](
    output: LayoutTensor[dtype, layout, MutAnyOrigin],
    input: LayoutTensor[dtype, layout, ImmutAnyOrigin],
):
    comptime assert (
        dtype.is_floating_point()
    ), "dtype must be a floating-point type"
    shared_max = LayoutTensor[
        dtype,
        Layout.row_major(BLOCK_DIM_X),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()
    shared_sum = LayoutTensor[
        dtype,
        Layout.row_major(BLOCK_DIM_X),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()
    global_i = Int(thread_idx.x)

    # Initialize out-of-bounds (shared_max[local_i], global_i >= input_size) shared memory addresses to the minimum
    # finite value for dtype, ensuring that if these elements are accessed in the parallel max reduction below they
    # do not influence the result (max(min_finite, x) == x for any x).
    var val: Scalar[dtype] = min_finite[dtype]()
    if global_i < input_size:
        val = rebind[Scalar[dtype]](input[global_i])
    shared_max[global_i] = val

    barrier()

    # Parallel reduction to find max similar to reduction we saw before
    stride = BLOCK_DIM_X // 2
    while stride > 0:
        if global_i < stride:
            shared_max[global_i] = max(
                shared_max[global_i], shared_max[global_i + stride]
            )
        barrier()
        stride = stride // 2

    block_max = shared_max[0]

    # Initialize out-of-bounds (shared_max[global_i], global_i >= input_size) shared memory addresses to 0.0,
    # ensuring that if these elements are accessed in the parallel sum reduction below they
    # do not influence the result (adding 0.0 does not change the sum).
    var exp_val: Scalar[dtype] = 0.0
    if global_i < input_size:
        exp_val = rebind[Scalar[dtype]](exp(val - block_max))
    shared_sum[global_i] = exp_val
    barrier()

    # Parallel reduction for sum similar to reduction we saw before
    stride = BLOCK_DIM_X // 2
    while stride > 0:
        if global_i < stride:
            shared_sum[global_i] += shared_sum[global_i + stride]
        barrier()
        stride = stride // 2

    block_sum = shared_sum[0]

    # Normalize by sum
    if global_i < input_size:
        output[global_i] = exp_val / block_sum

GPU 커널은 고도로 최적화된 병렬 리덕션 기법을 사용하여 수치적으로 안정적인 소프트맥스 알고리즘을 구현합니다. 커널을 상세히 분석해 보겠습니다:

커널 시그니처와 메모리 관리

fn softmax_gpu_kernel[
    layout: Layout,
    input_size: Int,
    dtype: DType = DType.float32,
](
    output: LayoutTensor[mut=True, dtype, layout],
    input: LayoutTensor[mut=False, dtype, layout],
)

커널의 파라미터 구성:

• 입출력 텐서에 공통으로 사용되는 레이아웃 파라미터
• 정수 파라미터로 지정되는 벡터 크기
• 기본값이 float32인 설정 가능한 데이터 타입
• 연산 결과를 직접 저장하는 변경 가능한(mutable) 출력 텐서
• 변경 불가능한(mut=False) 입력 텐서

공유 메모리 할당

shared_max = LayoutTensor[dtype, Layout.row_major(BLOCK_DIM_X), MutAnyOrigin, address_space = AddressSpace.SHARED].stack_allocation()
shared_sum = LayoutTensor[dtype, Layout.row_major(BLOCK_DIM_X), MutAnyOrigin, address_space = AddressSpace.SHARED].stack_allocation()

커널은 두 개의 공유 메모리 버퍼를 할당합니다:

• shared_max: 병렬 최댓값 탐색 리덕션용
• shared_sum: 병렬 합계 연산용
• 둘 다 BLOCK_DIM_X = 128 크기를 사용
• 공유 메모리는 블록 내 모든 스레드에 빠른 접근을 제공

스레드 인덱싱

global_i = thread_idx.x

이 소프트맥스 구현은 단일 1D 스레드 블록에서 동작합니다. 즉, 전역 인덱스와 로컬 인덱스가 동일합니다.

최댓값 탐색 단계

var val: Scalar[dtype] = min_finite[dtype]()
if global_i < input_size:
    val = rebind[Scalar[dtype]](input[global_i])

shared_max[local_i] = val
barrier()

각 스레드를 다음과 같이 초기화합니다:

• 유효 범위를 벗어난 요소에는 최소 유한(finite) 값 할당
• 유효한 요소에 매핑되는 스레드에는 실제 입력값 할당
• 리덕션 과정을 위해 공유 메모리에 저장
• 모든 스레드의 메모리 쓰기가 완료되도록 배리어 동기화

병렬 max 리덕션

stride = BLOCK_DIM_X // 2
while stride > 0:
    if local_i < stride:
        shared_max[local_i] = max(shared_max[local_i], shared_max[local_i + stride])
    barrier()
    stride = stride // 2

병렬 트리 리덕션 패턴을 구현합니다:

1. stride = 64(BLOCK_DIM_X의 절반)로 시작
2. 각 활성 스레드가 stride만큼 떨어진 두 값 비교
3. 더 작은 인덱스에 최댓값 저장
4. 배리어로 모든 스레드 동기화
5. Stride를 절반으로 줄이고 반복
6. \(\log_2(BLOCK\_DIM\_X)~\) 단계 후 shared_max[0]에 전체 최댓값이 담김

이 로그 리덕션은 대규모 입력에서 선형 스캔보다 훨씬 빠릅니다.

수치적으로 안정적인 지수 함수 적용

block_max = shared_max[0]

var exp_val: Scalar[dtype] = 0.0
if global_i < input_size:
    exp_val = rebind[Scalar[dtype]](exp(val - block_max))

각 스레드가 수행하는 작업:

1. 공유 메모리에서 전체 최댓값 읽음
2. 지수 함수를 적용하기 전에 입력값에서 최댓값 차감
3. 이 차감이 수치 안정성의 핵심 — 오버플로우 방지
4. 가장 큰 지수가 \(e^0 = 1\)이 되고, 나머지는 모두 \(e^{음수} < 1\)

병렬 sum 리덕션

shared_sum[local_i] = exp_val
barrier()

stride = BLOCK_DIM_X // 2
while stride > 0:
    if local_i < stride:
        shared_sum[local_i] += shared_sum[local_i + stride]
    barrier()
    stride = stride // 2

두 번째 리덕션 단계:

1. 모든 지수 값을 공유 메모리에 저장
2. max와 동일한 트리 기반 리덕션 패턴 사용
3. 단, 최댓값 비교 대신 덧셈 수행
4. \(\log_2(BLOCK\_DIM\_X)~\) 단계 후 shared_sum[0]에 모든 지수 값의 총합이 담김

최종 정규화

block_sum = shared_sum[0]

if global_i < input_size:
    output[global_i] = exp_val / block_sum

각 스레드가 수행하는 작업:

1. 공유 메모리에서 총합을 읽음
2. 자신의 지수 값을 이 총합으로 나눔
3. 정규화된 확률을 출력 버퍼에 기록
4. 합이 1인 유효한 확률 분포 생성

성능 특성

이 구현은 뛰어난 성능 특성을 갖습니다:

• 복잡도: 순차적 접근의 \(O(n)\)에 비해 max와 sum 계산 모두 \(O(\log n)\)
• 메모리 효율: 공유 메모리를 \(2 \times BLOCK\_DIM\_X~\) 요소만 사용
• 작업 효율: 각 스레드가 약 \(2 \times \log_2(BLOCK\_DIM\_X)~\) 회 연산 수행
• 부하 분산: 각 스레드가 동일한 양의 작업 처리
• 동기화: 필요한 곳에서만 최소한의 배리어 사용
• 메모리 접근: 최적 대역폭을 위한 병합된 전역 메모리 접근 패턴

이 알고리즘은 수치적으로도 견고합니다. 최댓값 차감 기법을 적용하여 신경망 활성화에서 흔한 넓은 범위의 값에서도 정밀도를 유지하며, 오버플로우/언더플로우 가능성을 처리합니다.

fn softmax_cpu_kernel[
    layout: Layout,
    input_size: Int,
    dtype: DType = DType.float32,
](
    output: LayoutTensor[dtype, layout, MutAnyOrigin],
    input: LayoutTensor[dtype, layout, ImmutAnyOrigin],
):
    comptime assert (
        dtype.is_floating_point()
    ), "dtype must be a floating-point type"
    var max_val: Scalar[dtype] = min_finite[dtype]()
    for i in range(input_size):
        max_val = max(max_val, rebind[Scalar[dtype]](input[i]))

    var sum_exp: Scalar[dtype] = 0.0
    for i in range(input_size):
        var exp_val = rebind[Scalar[dtype]](exp(input[i] - max_val))
        output[i] = exp_val
        sum_exp += exp_val

    for i in range(input_size):
        output[i] = output[i] / sum_exp

CPU 구현은 같은 수학적 접근 방식을 따르되 단일 스레드 실행에 최적화된 순차적 폴백을 제공합니다. 각 단계를 분석해 보겠습니다:

1. 최댓값 탐색:

   var max_val: Scalar[dtype] = min_finite[dtype]()
   for i in range(input_size):
       max_val = max(max_val, rebind[Scalar[dtype]](input[i]))
   

   최소 유한값으로 초기화하고 배열을 선형 스캔하며 만난 최댓값을 추적합니다. \(O(n)\) 복잡도이지만, 병렬화할 코어가 많지 않은 CPU에서는 효율적으로 동작합니다.

2. 지수 함수 적용과 합산:

   var sum_exp: Scalar[dtype] = 0.0
   for i in range(input_size):
       var exp_val = rebind[Scalar[dtype]](exp(input[i] - max_val))
       output[i] = exp_val
       sum_exp += exp_val
   

   각 요소에 대해 \(e^{x_i - max}\)를 계산하고 결과를 출력 버퍼에 저장하면서 합계 \(\sum_{j=1}^{n} e^{x_j - max}\)를 한 번의 순회로 누적합니다. 별도의 반복문을 사용하는 것에 비해 메모리 연산을 최소화합니다.

3. 정규화:

   for i in range(input_size):
       output[i] = output[i] / sum_exp
   

   마지막으로 각 요소를 합계로 나눠 소프트맥스 공식에 따른 올바른 확률 분포를 생성합니다:

   $$\Large \text{softmax}(x_i) = \frac{e^{x_i - \max(x)}}{\sum_{j=1}^{n} e^{x_j - \max(x)}}$$

CPU 구현은 동일한 수치 안정성 기법(최댓값 차감)을 사용하되, 병렬이 아닌 순차적 연산으로 처리합니다. 공유 메모리나 스레드 동기화를 다룰 필요가 없어 GPU 버전보다 단순하지만, 대규모 입력에서는 효율이 떨어집니다.

두 구현 모두 @compiler.register("softmax") 데코레이터를 통해 MAX 그래프의 커스텀 연산 시스템에 등록되므로, 가용 여부에 따라 어느 디바이스에서든 매끄럽게 실행됩니다.

    with Graph(
        "softmax_graph",
        input_types=[
            TensorType(
                dtype,
                shape=input_tensor.shape,
                device=DeviceRef.from_device(device),
            ),
        ],
        custom_extensions=[mojo_kernels],
    ) as graph:
        input_value = graph.inputs[0]

        # The output shape is the same as the input for softmax
        # Note: the name must match the name used in `@compiler.register("softmax")` in op/softmax.mojo
        output = ops.custom(
            name="softmax",
            values=[input_value],
            device=DeviceRef.from_device(device),
            out_types=[
                TensorType(
                    dtype=input_value.tensor.dtype,
                    shape=input_value.tensor.shape,
                    device=DeviceRef.from_device(device),
                )
            ],
            parameters={
                "target": "gpu" if device == Accelerator() else "cpu",
                "input_size": input_tensor.shape[0],
                "dtype": dtype,
            },
        )[0].tensor
        graph.output(output)

파이썬 통합은 NumPy 배열과 최적화된 Mojo GPU 커널 사이에 매끄러운 다리를 만듭니다. 구현은 여러 핵심 구성 요소로 이뤄져 있습니다:

1. 그래프 설정과 구성:

   with Graph(
       "softmax_graph",
       input_types=[
           TensorType(
               dtype,
               shape=input_tensor.shape,
               device=DeviceRef.from_device(device),
           ),
       ],
       custom_extensions=[mojo_kernels],
   ) as graph:
   

   “softmax_graph“라는 이름의 연산 그래프를 생성합니다:

   • 적절한 dtype과 shape으로 입력 텐서 타입 정의
   • 텐서를 대상 디바이스(CPU 또는 GPU)에 매핑
   • 지정된 디렉토리에서 커스텀 Mojo 연산 로드
   • custom_extensions 파라미터가 Mojo 구현과의 연결 핵심

2. 커스텀 연산 구성:

   output = ops.custom(
       name="softmax",
       values=[input_value],
       out_types=[
           TensorType(
               dtype=input_value.tensor.dtype,
               shape=input_value.tensor.shape,
               device=DeviceRef.from_device(device),
           )
       ],
       parameters={
           "target": "gpu" if device == Accelerator() else "cpu",
           "input_size": input_tensor.shape[0],
           "dtype": dtype,
       },
   )[0].tensor
   

   커스텀 연산을 다음과 같이 설정합니다:

   • Mojo 코드의 @compiler.register("softmax")와 일치하는 이름
   • 리스트로 전달되는 입력 값
   • 입력 shape과 타입에 맞는 출력 타입 정의
   • 대상 디바이스, 벡터 크기, 데이터 타입을 포함한 커널 필수 파라미터
   • [0].tensor로 첫 번째 반환 요소에서 텐서 추출

3. 그래프 출력 정의:

   graph.output(output)
   

   연산의 결과를 그래프의 출력으로 등록합니다.

메인 스크립트는 다음과 같은 꼼꼼한 검증을 포함합니다:

• 랜덤 입력 데이터 생성: np.random.randn(INPUT_SIZE).astype(np.float32)
• SciPy로 기대 결과 계산: scipy_softmax(input_array)
• 수치 정확도 검증: np.testing.assert_allclose(..., rtol=1e-5)
• 출력이 유효한 확률 분포인지 확인: np.sum(result.to_numpy())

이 구현은 고성능 Mojo 커널과 파이썬의 과학 컴퓨팅 생태계를 통합하는 MAX 그래프의 강력한 역량을 보여주며, 효율성과 사용 편의성을 동시에 제공합니다.