[furiosa-opt v0.5.1] dma_gather_unscaled의 64 KiB SPM 요청에서 V1 scheduler 실패

한국어 (Korean) 일반
1

furiosa-opt v0.5.1 (3f99736deef4a1faedf0c2ac6813c82a2b1c8d31)에서 아래 코드가 BIN/EDF를 생성하지 못합니다.

use furiosa_opt_std::prelude::*;

axes![Slices = 16, Indices = 128, Width = 8];

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

#[device(chip = 1)]
pub fn gather_i8_64k(
    ctx: &mut Context,
    table: &HbmTensor<i8, Chip, m![2, Width]>,
    index: &HbmTensor<i32, Chip, m![Slices, 1 # 16, Indices]>,
) -> HbmTensor<i8, Chip, m![Slices, Indices, Width]> {
    type Slice = m![Slices, 1 # 16];

    let index: DmTensor<i32, Chip, Cluster, Slice, m![Indices]> = index.to_dm(&mut ctx.tdma);
    let values: DmTensor<i8, Chip, Cluster, Slice, m![Indices, Width]> = table.dma_gather_unscaled(&index);

    values.to_hbm(&mut ctx.tdma)
}

위 코드를 furiosa-opt-examples/src/spm_64k_repro.rs로 추가하고 src/lib.rs에서 export합니다.

pub mod spm_64k_repro;

다음 명령으로 재현됩니다.

cargo-furiosa-opt compile \
  spm_64k_repro::gather_i8_64k \
  -p furiosa-opt-examples \
  --message-format human

visa: V1 failed for all operator schedule heuristics:
DfsPostOrder: BeamSearch failed after the beam shrank to zero at step 3/169; most frequent terminal error (1/1): Insufficient resources for T5 := DmaCommand(T4): {Spm: AllocRequest { wait_targets: [], size: 65536 }}
SethiUllman: BeamSearch failed after the beam shrank to zero at step 3/169; most frequent terminal error (1/1): Insufficient resources for T5 := DmaCommand(T4): {Spm: AllocRequest { wait_targets: [], size: 65536 }}
  • non-empty BIN: 0/1
  • non-empty EDF: 0/1
  • Indices = 64로만 변경하면 컴파일됩니다.

이 mapping의 64 KiB SPM 요청이 지원 대상이면 scheduler 수정을 부탁드립니다. 지원 대상이 아니면 allocation 상한과 공식 분할 방법을 알려주세요!

2

안녕하세요.

지금 형태처럼 큰 index tensor에 unscaled gather를 쓰는 것은 SPM 크기 제약 때문에 컴파일 불가능한게 맞습니다. 권장하는 방법은 scaled gather이고, unscaled gather를 유지해야 한다면 index를 SPM에 들어가는 크기로 쪼개서 여러 번 실행하셔야 합니다.

1. 원인

unscaled gather는 index list를 SPM에 올려두고 읽습니다. per-PE SPM 의 크기는 4KB인데, 올려주신 코드의 index는 Slices * Indices = 16 * 128 = 2048 i32 = 8 KB가 필요합니다. Indices를 64로 줄이면 4 KB로 딱 맞기 때문에 통과했던 것입니다. 다음 릴리즈에서는 실패가 조금 더 친절하게 표시되도록 개선될 예정입니다.

2. 권장 방법: scaled gather

scaled gather는 index를 DRAM에서 읽으므로 SPM 제약이 없습니다. unscaled gather 는 gather를 실행하면서 동시에 index를 scale해주는 기능인데, index 크기가 많아지는 경우 상당히 느리기 때문에, 최선의 성능을 위해서는 index를 미리 byte offset으로 scale해 둔 scaled gather를 사용하는것을 권장하고 있습니다. unscaled gather는 index 크기가 작은 특수한 경우에만 최적화 용도로 사용하는걸 권장합니다.

index를 host에서 미리 byte offset으로 만들어 두실 수 있다면 그대로 dma_gather_scaled에 넣으시면 되고, raw row index를 그대로 받아야 한다면 아래처럼 device에서 scale할 수 있습니다.

#[device(chip = 1)]
pub fn gather_scaled_from_raw_index(
    ctx: &mut Context,
    table: &HbmTensor<bf16, Chip, m![Rows, Width]>,
    index: &HbmTensor<i32, Chip, m![Slices, Indices]>,
) -> HbmTensor<bf16, Chip, m![Slices, Indices, Width]> {
    type Slice = m![Slices, Indices / 8];
    type Packet = m![Indices % 8];
    const ROW_BYTES: i32 = (<m![Width]>::SIZE * size_of::<bf16>()) as i32;

    let raw: DmTensor<i32, Chip, Cluster, Slice, Packet> = index.to_dm(&mut ctx.tdma);
    let scaled: DmTensor<i32, Chip, Cluster, Slice, Packet> = ctx
        .main
        .begin(raw.view())
        .fetch::<m![1], Packet>()
        .fetch_cast::<i32>()
        .collect::<m![1], m![Indices % 8]>()
        .vector_init()
        .vector_intra_slice_tag(TagMode::Zero)
        .vector_fxp(FxpBinaryOp::MulInt, ROW_BYTES)
        .vector_final()
        .commit_trim::<Packet>()
        .commit();
    let offsets: HbmTensor<i32, Chip, m![Slices, Indices]> = scaled.to_hbm(&mut ctx.tdma);

    let values: DmTensor<bf16, Chip, Cluster, Slice, m![Indices % 8, Width]> =
        table.dma_gather_scaled(&offsets);

    values.to_hbm(&mut ctx.tdma)
}

주의하실 점은 scale 값이 element 개수가 아니라 byte stride라는 것입니다. 위 예제의 payload는 bf16이라 Width = 8 elements = 16 bytes 이므로 8이 아닌 16을 곱해야 합니다.

3. unscaled gather를 유지해야 한다면: index 분할

문의하신 "공식 분할 방법"에 해당하는 답변입니다. index를 SPM에 들어가는 chunk로 나눠 loop을 돌고, 각 chunk의 gather 결과를 출력 tensor의 자기 slot에 써 넣으면 됩니다. 아래는 Indices / 32 = 4 chunk로 나눈 예로, chunk 하나당 Slices * 32 = 512 i32 = 2 KB라 4 KB 안에 들어옵니다.

#[device(chip = 1)]
pub fn gather_unscaled_split_index(
    ctx: &mut Context,
    table: &HbmTensor<bf16, Chip, m![Rows, Width]>,
    index: &HbmTensor<i32, Chip, m![Slices, Indices]>,
) -> HbmTensor<bf16, Chip, m![Slices, Indices, Width]> {
    type Slice = m![Slices, Indices % 32 / 2];
    type Packet = m![Indices % 2];

    let mut out = unsafe {
        HbmTensor::<bf16, Chip, m![Slices, Indices, Width]>::from_addr(0x0080_0000)
    };
    for c in 0..4 {
        let chunk: DmTensor<i32, Chip, Cluster, Slice, Packet> = index
            .view()
            .tile::<m![Indices / 32], 1, m![Slices, 1 # 4, Indices % 32]>(c)
            .to_dm(&mut ctx.tdma);
        let values: DmTensor<bf16, Chip, Cluster, Slice, m![Indices % 2, Width]> =
            table.dma_gather_unscaled(&chunk);

        values.view().to_hbm_view(
            &mut ctx.tdma,
            out.view_mut()
                .tile::<m![Indices / 32], 1, m![Slices, 1 #{!} 4, Indices % 32, Width]>(c),
        );
    }

    out
}

감사합니다.