Contraction Engine

Contraction Engine은 matmul, convolution 같은 이항 텐서 축약을 수행합니다. 빠른 시작에서 다룬 내용을 다시 떠올려 보십시오:

• 텐서 축약은 두 입력 텐서를 받아 두 텐서가 공유하는(축약되는) 축을 따라 리덕션합니다. 내적, GEMV, GEMM이 대표적인 예입니다.
• 축약은 브로드캐스트, 곱셈, 리덕션의 세 단계로 분해됩니다.
• 한 피연산자는 Collect Engine에서 스트리밍되고, 다른 피연산자는 TRF(Tensor Register File)에 저장되어 있습니다.
• 축약은 메인 컨텍스트에서 실행됩니다. TRF 준비는 .to_trf()를 통해 sub 컨텍스트에서 실행됩니다.

아키텍처

네 개의 파이프라인 단계가 워크로드를 나누어 맡습니다. 하나는 브로드캐스트와 곱셈을, 세 개는 리덕션을 담당합니다. 각 단계는 다른 단계와 겹치지 않는 자신만의 차원을 처리합니다.

%%{init: {'flowchart': {'htmlLabels': true}, 'themeCSS': '.cluster-label .nodeLabel { font-size: 16px; font-weight: 600; }'}}%%
flowchart TB
    CO[Collect Engine] --> SA
    TRF[(TRF)] --> TS

    subgraph CE[Contraction Engine]
        direction TB
        subgraph BC[Outer]
            direction LR
            SA[Stream Adapter]
            TS[TRF Sequencer]
            MUL[Elementwise Multiply]
            SA --> MUL
            TS --> MUL
        end
        SC[Packet Reducer]
        TR[Time Reducer]
        RR[Lane Folder]
        MUL --> SC
        SC --> TR
        TR --> RR
    end

    RR --> VE[Vector Engine]

    click SA "./outer.html#stream-adapter" "Stream Adapter"
    click TS "./outer.html#trf-sequencer" "TRF Sequencer"
    click SC "./packet-reducer.html" "Packet Reducer"
    click TR "./time-reducer.html" "Time Reducer"
    click RR "./lane-folder.html" "Lane Folder"
    click CO "../collect-engine.html" "Collect Engine"
    click TRF "../register-files.html#tensor-register-file" "Tensor Register File"
    click VE "../vector-engine/index.html" "Vector Engine"

• Outer (브로드캐스트와 곱셈): 두 피연산자를 일치하는 형태 [Chip, Cluster, Slice, Lane, Time, Packet]로 브로드캐스트하고 요소별로 곱하여 하나의 곱 텐서를 만듭니다. Chip / Cluster / Slice는 그대로 통과합니다. Lane은 TRF와 후속 리듀서가 공유하는 공간 병렬성을 인덱싱합니다. Time과 Packet은 함께 패킷 스트림을 나타냅니다. 세 개의 하위 단계가 직렬로 실행됩니다. Stream Adapter는 스트리밍 피연산자를 브로드캐스트하고, TRF Sequencer는 TRF 피연산자를 브로드캐스트하며, Multiplier는 축약 출력 타입으로 확장한 뒤(i4/i8 -> i32, f8/bf16 -> f32) 두 피연산자를 요소별로 곱합니다.
• Packet Reducer (Packet 내부 리덕션): Packet에 매핑된 축약 축을 따라 병렬 트리로 리덕션하며, 레인마다 트리가 하나씩 있습니다.
• Time Reducer (Time에 걸친 리덕션): 사이클별 결과를 공유 누산기에 누적합니다.
• Lane Folder (Lane 폴딩): 버퍼를 출력 스트림으로 내보내며, 모드에 따라 Lane을 OutPacket 또는 OutTime에 흡수시킵니다. Slice 간 또는 Chip 간 리덕션은 후속 단계에서 Vector Engine이 처리합니다.

Outer 단계는 (RNGD에서) Lane ≤ 8 및 Packet ≤ 64 B로 제한됩니다. 자세한 내용은 Packet Reducer와 Time Reducer를 참고하십시오.

예제: Batched MatMul

빠른 시작에서는 내적, GEMV, GEMM을 차례로 살펴봅니다. 이 절은 batched matmul 변형에 대한 아키텍처 레퍼런스로서 실행 가능한 변형 코드 조각을 모두 한곳에 모아 둡니다. 엔드투엔드 입문 경로는 여전히 빠른 시작입니다. Batched matmul은 GEMM에 선행 배치 축 V를 추가하여 확장한 것입니다: \(VMK, KN \rightarrow VMN\). 커널은 V개의 독립적인 (M × K) 입력 각각과 공유 (K × N) 가중치로부터 (M × N) 곱을 생성합니다.

아래의 세 가지 변형은 어떤 축이 Time에 놓이는지에 따라 커널을 분류합니다. 나머지 축은 공간 병렬성으로 활용됩니다. 세 변형은 다음 축을 공유합니다:

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![V = 32, M = 32, N = 8, K = 32];   // V batch, M×N output, K contraction
}

Time에 놓인 K

K(축약 축)는 Time에 놓입니다. M은 Cluster와 Slice에 걸쳐 분할되고 V도 분할됩니다. V % 16은 Slice에 합류하고 V / 16 = 2는 K와 함께 Time에 합류합니다(V × M = 1024는 RNGD의 Chip당 공간 셀 512개에 들어가지 않으므로 V의 바깥쪽 청크는 반복되어야 합니다). Packet은 1 # 32로 패딩되고, 리덕션은 Packet Reducer의 공간 트리를 거치지 않고 사이클에 걸쳐 순차적으로 진행됩니다. 따라서 사이클마다 32개의 곱셈기 중 1개만 유효한 작업을 수행합니다(bf16 기준 MAC 활용률 1/32). 결과적으로 퇴화된 커널이 되며, 이 커널은 교육용 기준선으로만 제시합니다.

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![V = 32, M = 32, N = 8, K = 32];   // V batch, M×N output, K contraction

type Chip    = m![1];                   // single chip
type Cluster = m![M / 16];              // outer M split across clusters (M / 16 = 2)
type Slice   = m![M % 16, V % 16];      // inner M × inner V = 16 × 16 = 256 slices per cluster
type Lane    = m![N];                   // N (output channels) partitions the 8 hardware lanes

/// Batched matmul with K placed in Time.
fn bmatmul_k_in_time<'l, const T: Tu>(
    // Streaming operand: V outer + K in Time, with a one-element Packet m![1].
    input: CollectTensor<'l, T, bf16, Chip, Cluster, Slice, m![V / 16, K], m![1 # 16]>,
    // TRF operand: N in Lane, K in Element. Stored into TRF by a prior .to_trf() call.
    trf: &TrfTensor<bf16, Chip, Cluster, Slice, Lane, m![K]>,
    // Output: one (M × N) f32 matrix per (slice, V-outer) pair.
) -> ContractTensor<'l, T, f32, Chip, Cluster, Slice, m![V / 16], m![N]> {
    input
         // Outer: Lane = m![N] (inferred from trf), OutTime = m![V / 16, K], OutPacket = m![1 # 32].
         // input: 1 K-element broadcast across all N lanes.
         // trf:   1 K-element per lane, advancing one K-step per cycle.
         .contract_outer::<m![V / 16, K], m![1 # 16], _, _, _>(trf)
         // Packet Reducer: OutPacket = m![1]. Nothing to reduce.
         .contract_packet::<m![1]>()
         // Time Reducer: OutTime = m![V / 16]. K iterates over Time and accumulates; V outer survives.
         .contract_time::<m![V / 16]>()
         // Lane Folder: Lane folds into OutPacket. Interleaved mode emits 8 lanes per cycle.
         .contract_lane::<m![V / 16], m![N]>(LaneMode::Interleaved)
}

let mut device = Device::new(Topology { chips: 1, pes: 8 }).unwrap();

let a: CollectTensor<'_, _, bf16, Chip, Cluster, Slice, m![V / 16, K], m![1 # 16]> = CollectTensor::new(&mut device.main, Tensor::zero());
let b: TrfTensor<bf16, Chip, Cluster, Slice, Lane, m![K]> = TrfTensor::zero();
let _o = bmatmul_k_in_time(a, &b);
}

이러한 병리적 사례를 피하려면 K를 Packet에 유지하고(Packet Reducer의 트리를 통한 병렬 리덕션) 남은 축(V, M, N)을 Cluster, Slice, Lane에 분산하여 공간 병렬성을 극대화하십시오. 아래의 두 전략은 이 원칙을 적용하며, 단순화를 위해 클래스마다 축을 하나씩 배치합니다. 실제 커널에서는 크기상 필요한 경우 하나의 축을 여러 클래스에 걸쳐 분할할 수 있습니다.

M을 Time에 배치

V(배치)는 Cluster와 Slice에 분산됩니다(Slice당 배치 요소 하나). M은 Time에, K는 Packet에 배치합니다. 이 전략은 (1) Slice 수가 배치를 감당할 수 있고, (2) N이 Lane에 들어가며, (3) K가 단일 Packet에 들어갈 때 적용할 수 있습니다. K가 Packet보다 크면 K를 Packet(공간)과 Time(시간)에 걸쳐 분할하십시오. 이 전략은 Lane 전반의 MAC 활용도를 극대화합니다.

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![V = 32, M = 32, N = 8, K = 32];   // V batch, M×N output, K contraction

type Chip    = m![1];                   // single chip
type Cluster = m![V / 16];              // outer V split across clusters (V / 16 = 2)
type Slice   = m![V % 16 # 256];        // inner V split across slices (V % 16 = 16 per cluster)
type Lane    = m![N];                   // N (output channels) partitions the 8 hardware lanes (N = 8 fills the cap)

/// Batched matmul: V slices × (M × K) · (K × N) → V × M × N.
fn bmatmul_m_in_time<'l, const T: Tu>(
    // Streaming operand: M in Time, K in Packet.
    // Element type can be i4, i8, f8, or bf16; integers widen to i32 output, floats to f32.
    input: CollectTensor<'l, T, bf16, Chip, Cluster, Slice, m![M, K / 16], m![K % 16]>,
    // TRF operand: N in Lane (one output channel per lane), K in Element.
    // Stored into TRF by a prior .to_trf() call in the sub context.
    trf: &TrfTensor<bf16, Chip, Cluster, Slice, Lane, m![K]>,
    // Output: one (M × N) f32 matrix per slice.
) -> ContractTensor<'l, T, f32, Chip, Cluster, Slice, m![M], m![N]> {
    input
         // Outer: broadcast input and trf, multiply elementwise.
         // Lane = m![N] (inferred from trf), OutTime = m![M], OutPacket = m![K].
         // input: K elements broadcast across all N lanes.
         // trf:   K elements per lane, broadcast across all M cycles.
         .contract_outer::<m![M], m![K], _, _, _>(trf)
         // Packet Reducer: OutPacket = m![1]. Sum K spatially via the reduction tree.
         .contract_packet::<m![1]>()
         // Time Reducer: OutTime = m![M]. Nothing to reduce.
         .contract_time::<m![M]>()
         // Lane Folder: Lane folds into OutPacket. Interleaved mode emits 8 lanes per cycle.
         .contract_lane::<m![M], m![N]>(LaneMode::Interleaved)
}

let mut device = Device::new(Topology { chips: 1, pes: 8 }).unwrap();

let a: CollectTensor<'_, _, bf16, Chip, Cluster, Slice, m![M, K / 16], m![K % 16]> = CollectTensor::new(&mut device.main, Tensor::zero());
let b: TrfTensor<bf16, Chip, Cluster, Slice, Lane, m![K]> = TrfTensor::zero();
let _o = bmatmul_m_in_time(a, &b);
}

V를 Time에 배치

V(배치)는 Time에, K는 Packet에 배치합니다. M은 Cluster와 Slice에 걸쳐 분할됩니다. 이 전략은 (1) Slice 수가 M을 감당할 수 있고(M / 16은 Cluster에, M % 16은 Slice에), (2) N이 Lane에 들어가며, (3) K가 단일 Packet에 들어갈 때 적용할 수 있습니다. 배치가 지배적인 축일 때(예: 배치 추론) 유용합니다.

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![V = 32, M = 32, N = 8, K = 32];   // V batch, M×N output, K contraction

type Chip    = m![1];                   // single chip
type Cluster = m![M / 16];              // outer M split across clusters (M / 16 = 2)
type Slice   = m![M % 16 # 256];        // inner M split across slices (M % 16 = 16 per cluster)
type Lane    = m![N];                   // N (output channels) partitions the 8 hardware lanes

/// Batched matmul with V (batch) placed in Time.
fn bmatmul_v_in_time<'l, const T: Tu>(
    // Streaming operand: V in Time, K in Packet.
    input: CollectTensor<'l, T, bf16, Chip, Cluster, Slice, m![V, K / 16], m![K % 16]>,
    // TRF operand: N in Lane, K in Element. Stored into TRF by a prior .to_trf() call.
    trf: &TrfTensor<bf16, Chip, Cluster, Slice, Lane, m![K]>,
    // Output: one (V × N) f32 matrix per slice.
) -> ContractTensor<'l, T, f32, Chip, Cluster, Slice, m![V], m![N]> {
    input
         // Outer: Lane = m![N] (inferred from trf), OutTime = m![V], OutPacket = m![K].
         // input: K elements broadcast across all N lanes.
         // trf:   K elements per lane, broadcast across all V cycles.
         .contract_outer::<m![V], m![K], _, _, _>(trf)
         // Packet Reducer: OutPacket = m![1]. Sum K spatially via the reduction tree.
         .contract_packet::<m![1]>()
         // Time Reducer: OutTime = m![V]. Nothing to reduce.
         .contract_time::<m![V]>()
         // Lane Folder: Lane folds into OutPacket. Interleaved mode emits 8 lanes per cycle.
         .contract_lane::<m![V], m![N]>(LaneMode::Interleaved)
}

let mut device = Device::new(Topology { chips: 1, pes: 8 }).unwrap();

let a: CollectTensor<'_, _, bf16, Chip, Cluster, Slice, m![V, K / 16], m![K % 16]> = CollectTensor::new(&mut device.main, Tensor::zero());
let b: TrfTensor<bf16, Chip, Cluster, Slice, Lane, m![K]> = TrfTensor::zero();
let _o = bmatmul_v_in_time(a, &b);
}