타일링 버전

fn matmul_tiled[
    layout: Layout, size: UInt
](
    output: LayoutTensor[dtype, layout_tiled, MutAnyOrigin],
    a: LayoutTensor[dtype, layout_tiled, ImmutAnyOrigin],
    b: LayoutTensor[dtype, layout_tiled, ImmutAnyOrigin],
):
    local_row = thread_idx.y
    local_col = thread_idx.x
    tiled_row = block_idx.y * TPB + local_row
    tiled_col = block_idx.x * TPB + local_col

    a_shared = LayoutTensor[
        dtype,
        Layout.row_major(TPB, TPB),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()
    b_shared = LayoutTensor[
        dtype,
        Layout.row_major(TPB, TPB),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()

    var acc: output.element_type = 0

    # Iterate over tiles to compute matrix product
    @parameter
    for tile in range((size + TPB - 1) // TPB):
        # Load A tile - global row stays the same, col determined by tile
        if tiled_row < size and (tile * TPB + local_col) < size:
            a_shared[local_row, local_col] = a[
                tiled_row, tile * TPB + local_col
            ]

        # Load B tile - row determined by tile, global col stays the same
        if (tile * TPB + local_row) < size and tiled_col < size:
            b_shared[local_row, local_col] = b[
                tile * TPB + local_row, tiled_col
            ]

        barrier()

        # Matrix multiplication within the tile
        if tiled_row < size and tiled_col < size:

            @parameter
            for k in range(min(Int(TPB), Int(size - tile * TPB))):
                acc += a_shared[local_row, k] * b_shared[k, local_col]

        barrier()

    # Write out final result
    if tiled_row < size and tiled_col < size:
        output[tiled_row, tiled_col] = acc

타일링 행렬 곱셈 구현은 작은 타일 \((3 \times 3)\) 을 사용하여 큰 행렬 \((9 \times 9)\) 을 효율적으로 처리하는 방법을 보여줍니다. 동작 방식은 다음과 같습니다:

1. 공유 메모리 할당

   Input matrices (9×9) - (3×3) 타일링에 딱 맞는 크기:
   A = [0  1  2  3  4  5  6  7  8 ]    B = [0  2  4  6  8  10 12 14 16]
       [9  10 11 12 13 14 15 16 17]        [18 20 22 24 26 28 30 32 34]
       [18 19 20 21 22 23 24 25 26]        [36 38 40 42 44 46 48 50 52]
       [27 28 29 30 31 32 33 34 35]        [54 56 58 60 62 64 66 68 70]
       [36 37 38 39 40 41 42 43 44]        [72 74 76 78 80 82 84 86 88]
       [45 46 47 48 49 50 51 52 53]        [90 92 94 96 98 100 102 104 106]
       [54 55 56 57 58 59 60 61 62]        [108 110 112 114 116 118 120 122 124]
       [63 64 65 66 67 68 69 70 71]        [126 128 130 132 134 136 138 140 142]
       [72 73 74 75 76 77 78 79 80]        [144 146 148 150 152 154 156 158 160]
   
   블록당 공유 메모리 (3×3):
   a_shared[TPB, TPB]  b_shared[TPB, TPB]
   

2. 타일 처리 루프

   타일 수 = 9 // 3 = 3개 (나머지 없이 딱 나눠짐!)
   
   각 타일에 대해:
   1. A와 B에서 타일 로드
   2. 부분 곱 계산
   3. 레지스터에 누적
   

3. 메모리 로딩 패턴

   • \((9 \times 9)\) 이 딱 나눠지므로 경계 검사가 기술적으로는 불필요하지만, 방어적 프로그래밍과 다른 행렬 크기에도 대응할 수 있도록 포함합니다.

        # A 타일 로드 - 전역 행은 그대로, 열은 타일에 의해 결정
        if tiled_row < size and (tile * TPB + local_col) < size:
            a_shared[local_row, local_col] = a[
                tiled_row, tile * TPB + local_col
            ]
     
        # B 타일 로드 - 행은 타일에 의해 결정, 전역 열은 그대로
        if (tile * TPB + local_row) < size and tiled_col < size:
            b_shared[local_row, local_col] = b[
                tile * TPB + local_row, tiled_col
            ]
     

4. 타일 내 연산

   for k in range(min(TPB, size - tile * TPB)):
       acc += a_shared[local_row, k] * b_shared[k, local_col]
   

   • 공유 메모리 뱅크 충돌 회피:

     Bank Conflict Free (Good):        Bank Conflicts (Bad):
     Thread0: a_shared[0,k] b_shared[k,0]  Thread0: a_shared[k,0] b_shared[0,k]
     Thread1: a_shared[0,k] b_shared[k,1]  Thread1: a_shared[k,0] b_shared[1,k]
     Thread2: a_shared[0,k] b_shared[k,2]  Thread2: a_shared[k,0] b_shared[2,k]
     ↓                                     ↓
     서로 다른 뱅크에 병렬 접근             b_shared가 열 우선이었다면
     (a_shared는 broadcast)               같은 뱅크에 직렬 접근
     

     공유 메모리 뱅크 충돌 설명:

     • 왼쪽 (Good): b_shared[k,threadIdx.x]는 서로 다른 뱅크에 접근하고, a_shared[0,k]는 모든 스레드에 브로드캐스트 됩니다
     • 오른쪽 (Bad): b_shared가 열 우선이었다면 스레드들이 동시에 같은 뱅크에 접근하게 됩니다
     • 핵심: 이것은 전역 메모리 병합이 아닌 공유 메모리 접근 패턴에 관한 것입니다
     • 뱅크 구조: 공유 메모리는 32개 뱅크로 구성되어 있으며, 여러 스레드가 동시에 같은 뱅크의 다른 주소에 접근할 때 충돌이 발생합니다

5. 동기화 지점

   barrier() 호출 시점:
   1. 타일 로딩 후
   2. 타일 연산 후
   

주요 성능 특성:

• \((3 \times 3)\) 타일로 \((9 \times 9)\) 행렬 처리 (딱 맞는 크기!)
• 공유 메모리로 빠른 타일 접근
• 병합된 메모리 접근으로 전역 메모리 트랜잭션 최소화
• 뱅크 충돌을 피하도록 최적화된 공유 메모리 레이아웃과 접근 패턴

1. 결과 기록:

   if tiled_row < size and tiled_col < size:
      output[tiled_row, tiled_col] = acc
   

   • 다른 행렬 크기와 타일링 전략을 위한 방어적 경계 검사 포함
   • 출력 행렬에 직접 대입
   • 모든 스레드가 유효한 결과를 기록

주요 최적화

1. 레이아웃 최적화:

   • 모든 텐서에 행 우선 레이아웃
   • 효율적인 2D 인덱싱

2. 메모리 접근:

   • 병합된 전역 메모리 로드
   • 효율적인 공유 메모리 활용

3. 연산:

   • 레지스터 기반 누적, 즉 var acc: output.element_type = 0
   • @parameter를 통한 컴파일 타임 루프 전개

이 구현은 다음을 통해 높은 성능을 달성합니다:

• LayoutTensor를 활용한 효율적인 메모리 접근
• 최적의 타일링 전략
• 적절한 스레드 동기화
• 세심한 경계 처리

from gpu.memory import async_copy_wait_all
from layout.layout_tensor import copy_dram_to_sram_async

comptime NUM_THREADS = TPB * TPB
comptime BLOCK_DIM_COUNT = 2


fn matmul_idiomatic_tiled[
    layout: Layout, size: UInt
](
    output: LayoutTensor[dtype, layout_tiled, MutAnyOrigin],
    a: LayoutTensor[dtype, layout_tiled, ImmutAnyOrigin],
    b: LayoutTensor[dtype, layout_tiled, ImmutAnyOrigin],
):
    local_row = thread_idx.y
    local_col = thread_idx.x
    tiled_row = block_idx.y * TPB + local_row
    tiled_col = block_idx.x * TPB + local_col

    # Get the tile of the output matrix that this thread block is responsible for
    out_tile = output.tile[TPB, TPB](Int(block_idx.y), Int(block_idx.x))
    a_shared = LayoutTensor[
        dtype,
        Layout.row_major(TPB, TPB),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()
    b_shared = LayoutTensor[
        dtype,
        Layout.row_major(TPB, TPB),
        MutAnyOrigin,
        address_space = AddressSpace.SHARED,
    ].stack_allocation()

    var acc: output.element_type = 0

    comptime load_a_layout = Layout.row_major(1, TPB)  # Coalesced loading
    comptime load_b_layout = Layout.row_major(1, TPB)  # Coalesced loading
    # Note: Both matrices stored in same orientation for correct matrix multiplication
    # Transposed loading would be useful if B were pre-transposed in global memory

    @parameter
    for idx in range(size // TPB):  # Perfect division: 9 // 3 = 3 tiles
        # Get tiles from A and B matrices
        a_tile = a.tile[TPB, TPB](Int(block_idx.y), Int(idx))
        b_tile = b.tile[TPB, TPB](Int(idx), Int(block_idx.x))

        # Asynchronously copy tiles to shared memory with consistent orientation
        copy_dram_to_sram_async[
            thread_layout=load_a_layout,
            num_threads=NUM_THREADS,
            block_dim_count=BLOCK_DIM_COUNT,
        ](a_shared, a_tile)
        copy_dram_to_sram_async[
            thread_layout=load_b_layout,
            num_threads=NUM_THREADS,
            block_dim_count=BLOCK_DIM_COUNT,
        ](b_shared, b_tile)

        # Wait for all async copies to complete
        async_copy_wait_all()
        barrier()

        # Compute partial matrix multiplication for this tile
        @parameter
        for k in range(TPB):
            acc += a_shared[local_row, k] * b_shared[k, local_col]

        barrier()

    # Write final result to output tile
    if tiled_row < size and tiled_col < size:
        out_tile[local_row, local_col] = acc

관용적 타일링 행렬 곱셈은 Mojo의 LayoutTensor API와 비동기 메모리 연산을 활용하여 깔끔한 구현을 제공합니다.

핵심 포인트: 이 구현은 두 행렬 모두 병합 로딩을 사용하여 표준 A × B 행렬 곱셈을 수행합니다.

이 구현이 하는 것:

• 행렬 연산: 표준 \(A \times B\) 곱셈 (\(A \times B^T\) 가 아님)
• 로딩 패턴: 두 행렬 모두 Layout.row_major(1, TPB)로 병합 접근
• 연산: acc += a_shared[local_row, k] * b_shared[k, local_col]
• 데이터 레이아웃: 로딩 시 전치 없음 - 두 행렬을 같은 방향으로 로드

이 구현이 하지 않는 것:

• \(A \times B^T\) 곱셈을 수행하지 않음
• 전치 로딩 패턴을 사용하지 않음
• 복사 과정에서 데이터를 전치하지 않음

\((9 \times 9)\) 행렬 크기에서는 완벽한 타일링이 이루어져 모든 경계 검사가 불필요합니다:

1. LayoutTensor 타일 API

   out_tile = output.tile[TPB, TPB](block_idx.y, block_idx.x)
   a_tile = a.tile[TPB, TPB](block_idx.y, idx)
   b_tile = b.tile[TPB, TPB](idx, block_idx.x)
   

   수동 좌표 계산 없이 “(block_idx.y, block_idx.x) 위치의 타일을 가져온다“를 직접 표현합니다. 자세한 내용은 문서를 참고하세요.

2. 비동기 메모리 연산

   copy_dram_to_sram_async[
      thread_layout = load_a_layout,
      num_threads = NUM_THREADS,
      block_dim_count = BLOCK_DIM_COUNT
   ](a_shared,a_tile)
   copy_dram_to_sram_async[
      thread_layout = load_b_layout,
      num_threads = NUM_THREADS,
      block_dim_count = BLOCK_DIM_COUNT
   ](b_shared, b_tile)
   async_copy_wait_all()
   

   이 연산들은:

   • 레지스터를 우회하는 전용 복사 엔진을 사용하여 연산과 메모리 전송의 중첩을 가능하게 합니다 (copy_dram_to_sram_async 참고)
   • 최적의 메모리 접근 패턴을 위한 특화된 스레드 레이아웃을 사용합니다
   • 수동 메모리 초기화가 불필요합니다
   • 중요:
     • 표준 GPU 로드는 이미 비동기적입니다. 이 함수들은 더 나은 리소스 활용과 레지스터 우회를 제공합니다
     • copy_dram_to_sram_async는 기본적으로 1D 스레드 블록(block_dim.y == block_dim.z == 1)을 가정하며, 별도 지정이 없으면 스레드 블록의 모든 스레드가 복사에 참여합니다. 다음을 지정하여 이 동작을 변경할 수 있습니다:
       • block_dim_count: 스레드 블록의 차원 수 (2D 스레드 블록 THREADS_PER_BLOCK_TILED = (TPB, TPB)의 경우 2)
       • num_threads: 스레드 블록의 스레드 수 (TPB*TPB == 9)

3. 최적화된 메모리 접근 레이아웃

   comptime load_a_layout = Layout.row_major(1, TPB)    # 병합 로딩
   comptime load_b_layout = Layout.row_major(1, TPB)    # 병합 로딩
   # 참고: 표준 A × B 곱셈에서 두 행렬 모두 같은 레이아웃을 사용
   

   현재 구현의 메모리 접근 분석:

   두 행렬 모두 전역 메모리에서 병합 로딩을 위해 Layout.row_major(1, TPB)를 사용합니다:

   • load_a_layout: 스레드들이 협력하여 행렬 A 행의 연속 원소를 로드
   • load_b_layout: 스레드들이 협력하여 행렬 B 행의 연속 원소를 로드
   • 핵심: 스레드 레이아웃은 복사 시 스레드 간 협력 방식을 결정하며, 최종 데이터 레이아웃과는 별개입니다

   실제 연산 패턴 (A × B임을 증명):

   # 현재 구현의 실제 연산
   acc += a_shared[local_row, k] * b_shared[k, local_col]
   
   # 이것은 C[i,j] = Σ(A[i,k] * B[k,j])에 해당
   # 즉, 표준 행렬 곱셈 A × B
   

   두 행렬이 같은 병합 로딩 패턴을 사용하는 이유:

   전역 메모리에서 타일 로딩:
   - Matrix A 타일: 스레드들이 A[block_row, k], A[block_row, k+1], A[block_row, k+2]... 로드 (연속)
   - Matrix B 타일: 스레드들이 B[k, block_col], B[k, block_col+1], B[k, block_col+2]... 로드 (연속)
   
   Layout.row_major(1, TPB)로 두 패턴 모두 병합
   

   세 가지 별개의 메모리 고려사항:

   1. 전역→공유 병합: Layout.row_major(1, TPB)로 병합 전역 메모리 접근 보장
   2. 공유 메모리 연산: a_shared[local_row, k] * b_shared[k, local_col]로 뱅크 충돌 회피
   3. 행렬 연산: 연산 패턴이 A × B를 결정 (A × B^T가 아님)

4. 완벽한 타일링으로 경계 검사 불필요

   @parameter
   for idx in range(size // TPB):  # 나머지 없는 나눗셈: 9 // 3 = 3
   

   \((9 \times 9)\) 행렬과 \((3 \times 3)\) 타일에서는 모든 타일이 정확히 꽉 차기 때문에 경계 검사가 필요 없습니다!

5. 방어적 경계 검사를 포함한 깔끔한 타일 처리

   # 완벽한 타일링에서도 방어적 경계 검사 포함
   if tiled_row < size and tiled_col < size:
       out_tile[local_row, local_col] = acc
   

   \((9 \times 9)\) 의 완벽한 타일링에서는 이 경계 검사가 기술적으로 불필요하지만, 방어적 프로그래밍과 다른 행렬 크기와의 일관성을 위해 포함합니다.

성능 고려사항

관용적 구현은 타일링의 성능 이점을 유지하면서 더 깔끔한 추상화를 제공합니다:

1. 메모리 지역성: 타일링을 통해 공간적, 시간적 지역성을 활용
2. 병합 접근: 특화된 로드 레이아웃으로 병합 메모리 접근 패턴 보장
3. 연산-메모리 중첩: 비동기 메모리 연산을 통한 중첩 가능
4. 공유 메모리 효율: 불필요한 공유 메모리 초기화 없음
5. 레지스터 압력: 최적의 연산 처리량을 위한 누적 레지스터 사용

이 구현은 고수준 추상화로도 성능 저하 없이 복잡한 GPU 알고리즘을 표현할 수 있음을 보여줍니다. 고수준의 표현력과 저수준의 성능 제어를 결합하는 Mojo의 철학을 잘 보여주는 예시입니다.

수동 타일링과의 주요 차이점

관용적 접근 방식은 단순히 더 깔끔할 뿐 아니라, 특화된 메모리 레이아웃과 비동기 연산 덕분에 성능도 더 좋을 수 있습니다.

참고: 전치 로딩은 언제 유용할까?

현재 구현은 전치 로딩을 사용하지 않습니다. 이 섹션은 레이아웃 시스템으로 할 수 있는 것을 보여주기 위한 교육적 내용입니다.

현재 구현 요약:

• 두 행렬 모두 Layout.row_major(1, TPB) 사용
• 표준 A × B 곱셈 수행
• 복사 중 데이터 전치 없음

전치 로딩을 사용하는 교육적 시나리오:

이 퍼즐은 두 행렬 모두 표준 병합 로딩을 사용하지만, 레이아웃 시스템의 유연성은 다른 시나리오에서 강력한 최적화를 가능하게 합니다:

# 예시: A × B를 계산하기 위해 사전 전치된 행렬 B^T를 로드
# (현재 구현에서는 이렇게 하지 않음)
comptime load_b_layout = Layout.row_major(TPB, 1)   # B^T를 병합 접근으로 로드
comptime store_b_layout = Layout.row_major(1, TPB)  # 공유 메모리에 B로 저장
copy_dram_to_sram_async[src_thread_layout=load_b_layout, dst_thread_layout=store_b_layout](b_shared, b_tile)

전치 로딩의 활용 사례 (이 퍼즐에서는 사용하지 않음):

1. 이미 전치된 입력 행렬: \(B\) 가 전역 메모리에 전치 상태로 저장되어 있는 경우
2. 다른 알고리즘: \(A^T \times B\), \(A \times B^T\), 또는 \(A^T \times B^T\) 계산
3. 메모리 레이아웃 변환: 행 우선과 열 우선 레이아웃 간 변환
4. 별도 전치 연산 없이 로드: 필요한 방향으로 데이터를 직접 로드

핵심 구분:

• 현재 구현: 두 행렬 모두 표준 \(A \times B\) 곱셈에 Layout.row_major(1, TPB) 사용
• 전치 로딩 예시: 이미 전치된 데이터나 다른 행렬 연산을 처리할 때 다른 레이아웃 사용

이것은 Mojo의 철학을 보여줍니다: 일반적인 경우에 고수준 추상화를 유지하면서도, 필요할 때 저수준 제어를 제공합니다.

---

요약: 핵심 정리

관용적 타일링 구현이 실제로 하는 것:

1. 행렬 연산: 표준 A × B 곱셈
2. 메모리 로딩: 두 행렬 모두 Layout.row_major(1, TPB)로 병합 접근
3. 연산 패턴: acc += a_shared[local_row, k] * b_shared[k, local_col]
4. 데이터 레이아웃: 로딩 시 전치 없음

이것이 최적인 이유:

• 병합 전역 메모리 접근: Layout.row_major(1, TPB)로 효율적인 로딩 보장
• 뱅크 충돌 회피: 공유 메모리 접근 패턴이 충돌을 방지
• 표준 알고리즘: 가장 일반적인 행렬 곱셈 패턴을 구현