스케줄

스케줄러는 컨텍스트, 소스 순서, 메모리 위치를 사용하여 연산을 구성합니다. 스케줄러는 순차 실행과 동등한 결과를 얻는 데 필요한 의존성과 리소스 제약 조건을 보존합니다.

실행 컨텍스트

공개 실행 컨텍스트는 다음과 같습니다.

Main (device.main): 주 Tensor Unit 파이프라인을 구동합니다.
Sub (device.sub): register-file 프리로드 같은 연산이 사용하는 Tensor Unit 파이프라인의 일부를 구동합니다.
DMA (device.tdma): HBM, DM 및 기타 메모리 계층 간의 텐서 이동을 구동합니다.

하나의 컨텍스트 안에 있는 연산들은 순서가 정해져 있습니다. MNIST 스케줄에서 서로 다른 컨텍스트의 연산들은 스케줄링 리소스를 공유하지 않고 메모리 의존성을 위반하지 않는 한 겹쳐서 실행됩니다. 예를 들어, DMA 로드는 소비자가 로드된 값을 필요로 할 때까지 메인 컨텍스트 계산과 겹쳐 실행됩니다.

리소스 및 뱅크 경합

스케줄링 모델: 스케줄러는 컨텍스트 점유 정보를 사용합니다. 연산 A가 컨텍스트(예: 메인 컨텍스트)를 점유하면, 그 컨텍스트를 사용하는 다음 연산 B는 A가 완료될 때까지 대기합니다. 연산이 어떤 컨텍스트를 점유하는지 이해하면 병렬 실행을 예측할 수 있습니다.
컴파일러 스케줄링 동작: Tensor Unit 연산이 64회 접근 제한을 위반하게 되는 경우, 컴파일러는 해당 연산이 DMA를 점유하는 것처럼 스케줄링하여 동시 DMA 연산을 막습니다. 이는 데이터 준비와 계산이 병렬로 이루어지는 TCP 아키텍처 고유의 main/sub/DMA 컨텍스트 병렬성을 희생하지만, 치명적인 하드웨어 리셋을 방지합니다. 이를 엄격한 제약 조건으로 취급하십시오. 동일 뱅크에 64회 이상 연속으로 접근하는 패턴은 절대 사용하지 마십시오.
64회 접근 제한 세부 사항:
- 이 제한은 누적됩니다. 명령 간에 교차된 접근도 이 합계에 누적되므로, 모든 엔진이 동일 뱅크에 접근하는 총 횟수는 64 미만으로 유지되어야 합니다. 예를 들어, main: 30, sub: 20, DMA: 1은 합계 51(안전)이지만, main: 30, sub: 35, DMA: 1은 합계 66(기아 상태 유발)입니다.
- 컴파일러는 개별 명령 각각의 동일 뱅크 연속 접근을 64회 미만으로 유지하지만, 여러 명령이 동시에 실행될 때 합계가 64에 도달하는 것은 막을 수 없습니다.
- 실제로 서브 컨텍스트가 같은 뱅크에 연속으로 접근하는 경우는 드뭅니다(StoTrf, StoVrf 연산은 일반적으로 순차 주소를 사용하며, 타일링이 같은 뱅크 접근을 방지합니다).
- 한도를 초과하게 되는 서브 컨텍스트 연산 또한 DMA와 동시에 스케줄링되지 않습니다.
메인/서브 컨텍스트 경합: 메인 컨텍스트가 서브 컨텍스트를 기아 상태에 빠뜨릴 수 있지만, 이는 덜 심각합니다.
- DMA 기아와 달리, 서브 컨텍스트 기아는 NoC 타임아웃이나 하드웨어 리셋을 일으키지 않으며 처리 시간만 늘립니다.
- 충돌 확률이 더 낮습니다. DMA Engine은 한 번에 16개 뱅크를 점유하는 반면, sub fetch/commit 엔진은 하나의 뱅크만 점유합니다.
- 파이프라인 백프레셔로 인해 동일 컨텍스트 내의 fetch 엔진과 commit 엔진 사이에서는 기아 상태가 발생하지 않습니다.
- 성능 영향 예시: 서브 컨텍스트 stos 명령이 스케줄링된 동안 메인 컨텍스트 exec 명령이 특정 뱅크에 계속 접근하면, 서브 컨텍스트 처리가 지연됩니다. 최악의 경우: 전체 시간 = 메인 컨텍스트 시간 + 서브 컨텍스트 시간입니다. 이상적인 경우: 메인과 서브가 서로 다른 뱅크에 접근하여 전체 시간 = max(메인 컨텍스트 시간, 서브 컨텍스트 시간)을 달성합니다.

연산 순서와 의존성

스케줄러는 겹침을 드러내기 위해 독립적인 연산의 순서를 바꿀 수 있습니다. 순서를 바꾸면 값이 달라질 수 있는 경우에는 의존성 때문에 순서를 바꿀 수 없습니다. 소스 코드에 작성된 순서만 보고 성능이 개선되었다고 판단하지 마십시오. 생성된 스케줄을 살펴보고 튜닝 프로토콜에 따라 makespan을 비교하십시오.

Tensor Unit 연산은 HBM 텐서를 직접 사용할 수 없습니다. 데이터는 빠른 시작에서 설명한 메모리 계층을 거쳐 이동해야 합니다. 메모리 위치는 여러 컨텍스트가 공유하므로 스케줄러는 다음 해저드를 지켜야 합니다.

쓰기 후 읽기(Read after write): 소비자는 생산자의 쓰기를 기다립니다.
읽기 후 쓰기(Write after read): 뒤에 오는 쓰기는 앞선 읽기가 모두 끝날 때까지 기다립니다.
쓰기 후 쓰기(Write after write): 겹치는 쓰기는 프로그램 순서를 유지합니다.

연산 유형이 서로 다른 독립 컨텍스트도 때때로 기다리는 이유가 바로 이러한 해저드 때문입니다.

서로 다른 컨텍스트라도 같은 스케줄링 자원이 필요하면 여전히 직렬화됩니다. 예를 들어 메인 컨텍스트의 Vector Engine 연산과 서브 컨텍스트의 Vector Engine 연산은 해당 자원을 동시에 점유할 수 없습니다. 튜닝 수단을 고르기 전에 이러한 자원 대기와 데이터 의존성을 구별하여 진단해야 합니다.

MNIST 예제는 2계층 MLP의 스케줄 예제를 제공하며, 첫 번째 계층에서 ReLU를 적용합니다. 이 예제의 스케줄 이미지는 컨텍스트 간 겹침과 대기를 보여 줍니다.

flowchart LR
    X["Input<br>(X = 800)"]
    H["Hidden<br>(H = 256)"]
    C["Output<br>(C = 16)"]
    X -- "FC1 + ReLU" --> H
    H -- "FC2" --> C

디바이스 함수는 두 계층을 소스 코드 순서대로 호출합니다. 각 계층은 자신의 계산 경로에 필요한 메모리 계층에 피연산자를 준비합니다.

axes![X = 800, H = 256, C = 16];

type Chip = m![1];
type Cluster = m![1 # 2];

#[device(chip = 1)]
pub fn forward(
    device: &mut Device,
    input: &HbmTensor<bf16, Chip, m![X]>,
    fc1_weight: &HbmTensor<bf16, Chip, m![H, X]>,
    fc1_bias: &HbmTensor<bf16, Chip, m![H]>,
    fc2_weight: &HbmTensor<bf16, Chip, m![C, H]>,
    fc2_bias: &HbmTensor<bf16, Chip, m![C]>,
) -> HbmTensor<bf16, Chip, m![C]> {
    let hidden = fc1_relu(device, input, fc1_weight, fc1_bias);
    fc2(device, hidden, fc2_weight, fc2_bias)
}

각 완전 연결 계층은 행렬-벡터 곱을 계산하고 바이어스를 더합니다. 첫 번째 계층은 같은 패스에서 ReLU도 적용합니다.

전체 스케줄은 다음과 같습니다. 이어지는 이미지는 컨텍스트 동작과 의존성 동작을 확대하여 보여 줍니다.

MNIST full timeline