Collect Engine

모든 다운스트림 엔진(Contraction Engine, Vector Engine, Cast Engine, Transpose Engine, Commit Engine)은 정확히 32바이트 크기의 플릿(flit)을 소비합니다. Collect Engine은 임의 크기의 패킷을 두 단계를 거쳐 하나의 플릿으로 정규화합니다.

입력 패킷을 다음 32바이트 경계까지 패딩합니다. 패킷이 이미 32바이트로 정렬되어 있으면 이 단계를 건너뜁니다.
플릿 경계에서 분할합니다. 안쪽 32바이트는 Packet2가 되고, 바깥쪽 플릿 수는 Time2에 흡수됩니다. 패킷이 이미 32바이트이면 이 단계를 건너뜁니다.

그 결과로 생성된 CollectTensor는 파이프라인을 따라 다운스트림 엔진으로 전달되거나 Register File에 저장됩니다.

인터페이스

SwitchTensor와 FetchTensor는 모두 동일한 의미의 .collect()를 제공합니다. Slice 분배가 필요하지 않을 때 FetchTensor 진입점은 Switch Engine을 우회합니다.

impl<'l, const T: Tu, P: CanApplyCollect, D: Scalar, Chip: M, Cluster: M, Slice: M, Time: M, Packet: M, B: Backend>
    TuTensor<'l, T, P, D, Chip, Cluster, Slice, Time, Packet, B>
{
    /// Normalizes packet to exactly 32 bytes (one flit).
    ///
    /// Pads to flit-aligned boundary, then splits: inner 32 bytes become
    /// `Packet2`, outer flit portion is absorbed into `Time2`. For packets
    /// already ≤ 32 bytes, only padding is added.
    #[primitive(TuTensor::collect)]
    pub fn collect<Time2: M, Packet2: M>(self) -> CollectTensor<'l, T, D, Chip, Cluster, Slice, Time2, Packet2, B> {
        verify_collect::<D, Time, Packet, Time2, Packet2>();
        CollectTensor::new(self.device, self.inner.transpose(false))
    }
}

예제

단일 플릿 패킷

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 8, B = 32];

fn collect_identity<'l, const T: Tu>(
    input: SwitchTensor<'l, T, i8, m![1], m![1 # 2], m![1 # 256], m![A], m![B]>,
) -> CollectTensor<'l, T, i8, m![1], m![1 # 2], m![1 # 256], m![A], m![B # 32]> {
    // B=32 elements × 1 byte (i8) = 32 bytes = one flit.
    // Time and Packet pass through unchanged.
    input.collect()
}

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

let c: SwitchTensor<'_, _, i8, m![1], m![1 # 2], m![1 # 256], m![A], m![B]> = SwitchTensor::new(&mut device.main, Tensor::zero());
let _o = collect_identity(c);
}

입력 패킷이 이미 정확히 32바이트이면 collect는 패킷을 변경 없이 그대로 통과시킵니다(i8의 경우 B = 32개 요소 × 1바이트 = 32바이트).

Before:   Time = m![A]
          Packet = m![B]
          ┌──────────────────────────┐
          │            B             │  32 bytes
          └──────────────────────────┘

After:    Time = m![A]
          Packet = m![B # 32]
          ┌──────────────────────────┐
          │          B # 32          │  32 bytes
          └──────────────────────────┘

플릿 미만 패킷

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 8, B = 16];

fn collect_padding<'l, const T: Tu>(
    input: SwitchTensor<'l, T, i8, m![1], m![1 # 2], m![1 # 256], m![A], m![B]>,
) -> CollectTensor<'l, T, i8, m![1], m![1 # 2], m![1 # 256], m![A], m![B # 32]> {
    // B=16 elements × 1 byte = 16 bytes < 32 bytes.
    // Padded to 32 bytes: Packet2 = m![B # 32].
    // Time unchanged since it fits in one flit.
    input.collect()
}

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

let c: SwitchTensor<'_, _, i8, m![1], m![1 # 2], m![1 # 256], m![A], m![B]> = SwitchTensor::new(&mut device.main, Tensor::zero());
let _o = collect_padding(c);
}

입력 패킷이 32바이트보다 작으면 collect는 32바이트까지 패딩합니다(i8의 경우 B = 16개 요소 × 1바이트 = 16바이트).

Before:   Time = m![A]
          Packet = m![B]
          ┌────────────┐
          │     B      │  16 bytes
          └────────────┘

After:    Time = m![A]
          Packet = m![B # 32]
          ┌────────────┬─────────────┐
          │     B      │     pad     │  32 bytes
          └────────────┴─────────────┘

다중 플릿 패킷

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 8, B = 32];

fn collect_multi_flit<'l, const T: Tu>(
    input: SwitchTensor<'l, T, bf16, m![1], m![1 # 2], m![1 # 256], m![A], m![B]>,
) -> CollectTensor<'l, T, bf16, m![1], m![1 # 2], m![1 # 256], m![A, B / 16], m![B % 16]> {
    // B=32 elements × 2 bytes (bf16) = 64 bytes = 2 flits.
    // Inner 16 elements = 32 bytes → Packet2 = m![B % 16].
    // Outer 2 flits → absorbed into Time2 = m![A, B / 16].
    input.collect()
}

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

let c: SwitchTensor<'_, _, bf16, m![1], m![1 # 2], m![1 # 256], m![A], m![B]> = SwitchTensor::new(&mut device.main, Tensor::zero());
let _o = collect_multi_flit(c);
}

입력 패킷이 32바이트를 초과하면 collect는 패킷을 플릿 단위로 분할하고 바깥쪽 플릿 수를 Time에 흡수합니다(bf16의 경우 B = 32개 요소 × 2바이트 = 64바이트이므로 B / 16 = 2개 플릿).

Before:   Time = m![A]
          Packet = m![B]
          ┌──────────────────────────┬──────────────────────────┐
          │       B / 16 == 0        │       B / 16 == 1        │  64 bytes
          └──────────────────────────┴──────────────────────────┘
                    32 bytes                   32 bytes

After:    Time = m![A, B / 16]
          Packet = m![B % 16]
          ┌──────────────────────────┐
          │          B % 16          │  32 bytes  × B/16 time steps
          └──────────────────────────┘

패딩이 있는 다중 플릿 패킷

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 8, B = 56];

fn collect_multi_flit_padded<'l, const T: Tu>(
    input: SwitchTensor<'l, T, i8, m![1], m![1 # 2], m![1 # 256], m![A], m![B]>,
) -> CollectTensor<'l, T, i8, m![1], m![1 # 2], m![1 # 256], m![A, B # 64 / 32], m![B # 64 % 32]> {
    // B is not 32-byte aligned; first pad B to a multiple of 32 bytes.
    // B # 64=64 elements × 1 byte (i8) = 64 bytes = 2 flits.
    // Inner 32 elements = 32 bytes → Packet2 = m![B # 64 % 32].
    // Outer 2 flits → absorbed into Time2 = m![A, B # 64 / 32].
    input.collect()
}

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

let c: SwitchTensor<'_, _, i8, m![1], m![1 # 2], m![1 # 256], m![A], m![B]> = SwitchTensor::new(&mut device.main, Tensor::zero());
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { collect_multi_flit_padded(c) }));
}

입력 패킷이 32바이트로 정렬되어 있지 않으면 먼저 패딩됩니다(i8의 경우 B = 51개 요소 × 1바이트 = 51바이트이며, 64바이트로 패딩). 그런 다음 collect는 패킷을 플릿 단위로 분할하고 바깥쪽 플릿 수(B # 64 / 32 = 2)를 Time에 흡수합니다.

Before:   Time = m![A]
          Packet = m![B]
          ┌──────────────────────────┬───────────────┐
          │       B / 32 == 0        │  B / 32 == 1  │  51 bytes
          └──────────────────────────┴───────────────┘
                    32 bytes             19 bytes

Padded:   Time = m![A]
          Packet = m![B # 64]
          ┌──────────────────────────┬───────────────┬──────────┐
          │       B / 32 == 0        │  B / 32 == 1  │   pad    │  64 bytes
          └──────────────────────────┴───────────────┴──────────┘
                    32 bytes                   32 bytes

After:    Time = m![A, B # 64 / 32]
          Packet = m![B # 64 % 32]
          ┌──────────────────────────┐
          │       B # 64 % 32        │  32 bytes  × B # 64 / 32 time steps
          └──────────────────────────┘

Register File 로딩

정규화 후에는 .to_trf()로 CollectTensor를 Tensor Register File에 저장하거나 .to_vrf()로 Vector Register File에 저장하십시오. 해당 API, 영역, 용량, 주소 모드, 예제는 Register File 장에서 다룹니다.

Keyboard shortcuts

Tensor Contraction Processor 프로그래밍