Register File

Collect Engine은 Contraction Engine과 Vector Engine으로 데이터를 스트리밍합니다. 두 엔진은 Slice마다 하나씩 있는 Register File에서도 입력을 받습니다. Tensor Register File(TRF)은 Contraction Engine에 입력을 공급하고, Vector Register File(VRF)은 Vector Engine에 입력을 공급합니다. 이 Register File들은 이를 소비하는 엔진이 실행되기 전에 채워져 있어야 합니다.

Tensor Register File

인터페이스

TrfTensor는 TRF에 저장된 텐서입니다:

/// Tensor stored in the tensor register file.
#[primitive(TrfTensor)]
#[derive(Debug)]
pub struct TrfTensor<D: Scalar, Chip: M, Cluster: M, Slice: M, Lane: M, Element: M, B: Backend = CurrentBackend> {
    pub(crate) inner: Tensor<D, Pair<Chip, Pair<Cluster, Pair<Slice, Pair<Lane, Element>>>>, B>,
    _marker: PhantomData<(D, Chip, Cluster, Slice, Lane, Element)>,
}

Chip / Cluster / Slice는 소스에서 그대로 전달됩니다. Lane은 공간 병렬성(활성 레인 1, 2, 4 또는 8개)을 인덱싱합니다. Element는 레인별 레이아웃을 담습니다.

Collect Engine에서

CollectTensor에 .to_trf::<Lane, Element>()를 호출하면 TrfTensor가 생성됩니다. 데이터가 Register File의 어느 위치에 놓일지는 컴파일러가 결정하므로, 호출에서는 영역을 지정하지 않습니다:

impl<'l, const T: Tu, P: CanApplyToTrf, 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>
{
    /// Stores to the tensor register file.
    #[primitive(TuTensor::to_trf)]
    pub fn to_trf<Lane: M, Element: M>(self) -> TrfTensor<D, Chip, Cluster, Slice, Lane, Element, B> {
        verify_to_trf::<D, Lane, Time, Packet, Element>();
        TrfTensor::new(self.inner.transpose(false))
    }
}

.to_trf는 스트리밍되는 Time / Packet을 Lane / Element로 재구성합니다:

Lane    = Time / FlitsPerLane
Element = [Time % FlitsPerLane, Packet]

여기서 FlitsPerLane은 컴파일러가 Lane과 Time으로부터 도출하는 값이며, 각 레인은 FlitsPerLane개의 연속된 플릿으로 채워집니다.

다음의 간단한 사례는 TRF 전체를 점유합니다:

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

fn load_trf<'l, const T: Tu>(
    input: CollectTensor<'l, T, i8, m![1], m![1 # 2], m![1 # 256], m![1], m![B]>,
) -> TrfTensor<i8, m![1], m![1 # 2], m![1 # 256], m![1], m![B]> {
    input.to_trf()
}

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

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

예를 들어 matmul 커널에서 Lane은 출력 채널을 담고 Element는 축약되는 축을 담습니다.

#![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];

type Chip    = m![1];
type Cluster = m![V / 16];
type Slice   = m![V % 16 # 256];
type Lane    = m![N];

/// Stores matmul weights into TRF for consumption by `bmatmul` in
/// [Contraction Engine: Example: Batched MatMul](./contraction-engine/index.md#example-batched-matmul).
fn store_bmatmul_trf<'l, const T: Tu>(
    input: CollectTensor<'l, T, bf16, Chip, Cluster, Slice, m![N, K / 16], m![K % 16]>,
) -> TrfTensor<bf16, Chip, Cluster, Slice, Lane, m![K]> {
    input.to_trf()
}

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

let c: CollectTensor<'_, _, bf16, Chip, Cluster, Slice, m![N, K / 16], m![K % 16]> = CollectTensor::new(&mut device.main, Tensor::zero());
let _o = store_bmatmul_trf(c);
}

Vector Engine에서

저장이 반드시 collect에 이어서 이루어져야 하는 것은 아닙니다. 메인 컨텍스트 스트림은 Vector Engine을 먼저 실행한 뒤 그 패스가 계산한 결과를 저장할 수 있습니다:

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

fn store_pass(device: &mut Device) -> VrfTensor<f32, m![1], m![1 # 2], m![1 # 256], m![B]> {
    let dm: DmTensor<f32, m![1], m![1 # 2], m![1 # 256], m![B]> = DmTensor::new();
    device.main
        .begin(dm.view())
        .fetch::<m![1], m![B]>()
        .collect::<m![B / 8], m![B % 8]>()
        .vector_init()
        .vector_intra_slice_tag(TagMode::Zero)
        .vector_narrow_split::<m![B / 8, B % 8 / 4 % 2], m![B % 4]>()
        .vector_fp_unary(FpUnaryOp::Sqrt)
        .vector_widen_concat::<m![B / 8], m![B % 8]>()
        .vector_final()
        .to_vrf(&mut device.sub)
}

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

두 위치는 하나의 경로를 두 가지로 표기한 것이 아니라 서로 다른 데이터 경로입니다. collect에서 바로 저장하면 Fetch Network가 플릿을 Register File로 넘기며, 그 사이에 실행되는 엔진은 없습니다. vector_final에 이어 저장하면 패스가 Vector Engine 자체의 쓰기 포트에서 끝나므로, 레지스터에 놓이는 값은 패스가 계산한 값입니다. Element2는 패스 이후의 Time / Packet을 평탄화하며, 이는 도중에 reduce나 widen에 의해 형태가 바뀌었을 수 있습니다.

DM 직접 저장과 collect 위치 저장은 모두 Data Memory를 읽으므로, 계산된 값이 VRF에 도달하는 방법은 이것뿐입니다. 같은 명령에서 Vector Engine 이후에는 아무것도 실행될 수 없으므로, 저장을 cast, transpose 또는 commit과 결합할 수 없습니다.

저장이 점유하는 컨텍스트

main 컨텍스트에서 발행된 저장은 위의 두 위치 중 어느 쪽을 택하든 sub 컨텍스트를 통해 쓰기를 실행합니다. 따라서 두 컨텍스트를 모두 점유하며 &mut device.sub도 함께 받습니다. device.main과 device.sub는 별개의 필드이므로, &mut Device를 보유한 커널은 위 예제처럼 두 필드를 동시에 빌릴 수 있습니다.

빌림은 해당 문(statement)이 끝날 때까지 유지되며, 이것이 점유의 범위입니다. 명령이 완료되면 sub 컨텍스트는 다시 사용 가능해집니다. 레지스터가 얼마나 오래 점유되는지는 별개의 문제이며 답도 다릅니다(피연산자를 마지막으로 읽을 때까지). 이 할당은 컴파일러가 결정하므로, 반환된 VrfTensor는 빌림을 갖지 않습니다.

sub 컨텍스트에서 발행된 저장은 이미 해당 컨텍스트를 보유하고 있으므로 인수가 필요하지 않습니다.

Data Memory에서

완전히 연속적인 입력 접근(간격이나 재정렬 없음)의 경우, TRF는 단축 명령(StoTRF)을 지원합니다. 이는 전체 Fetch → Switch → Collect → to_trf() 파이프라인을 거치지 않고 Data Memory에서 TRF로 데이터를 직접 로드하는 간결한 하드웨어 명령입니다. 이 단축 경로는 임의 레이아웃 지원을 포기하는 대신 설정 오버헤드를 줄입니다.

Contraction Engine으로

각 읽기는 8개 레인 × (1개 또는 2개) 뱅크 × 1개 행 × 뱅크당 320비트를 다루며, 좁은 읽기(뱅크 1개)에서는 사이클당 레인마다 320비트를, 넓은 읽기(뱅크 2개)에서는 사이클당 레인마다 640비트를 만들어 냅니다. 모든 활성 레인은 같은 행의 뱅크 하나 또는 둘 모두에 병렬로 접근합니다(좁은 읽기는 뱅크 하나, 넓은 읽기는 두 뱅크 모두). 8개 lane 전체를 Slice 단위로 합산하면 사이클당 320바이트(좁은 읽기) 또는 사이클당 640바이트(넓은 읽기)입니다. Sequencer가 이러한 읽기를 row와 브로드캐스트에 걸쳐 어떻게 반복하는지는 TRF Sequencer를 참조하십시오.

아키텍처

TRF는 Slice당 8개 레인 × 2개 뱅크 × 128개 행 × 320비트 = 80 KB 구조를 가진 뱅크 방식 SRAM입니다. 8개 lane은 병렬로 동작하며, 접근마다 1개, 2개, 4개 또는 8개가 활성화됩니다.

320비트 row 하나에 몇 개의 요소가 들어가는지는 데이터 타입에 따라 달라집니다.

저장 시 승격되는 것은 정수뿐입니다. i4 → i5(5비트)와 i8 → i9(9비트)는 Fetch Adapter의 선택적 영점 차감을 위한 여유를 남겨 두며, 이 차감은 정수 가중치를 1비트 넓힙니다. i8 / f8과 bf16은 원래 너비(각각 8비트와 16비트)를 유지합니다. 320비트 행은 8비트 또는 16비트를 빈틈없이 패킹하는 경우보다 여유 공간을 더 가지므로, 같은 물리적 행 너비로 모든 타입을 처리할 수 있습니다.

활성 lane이 8개보다 적으면 각 활성 lane은 더 많은 row를 보게 됩니다(마치 row 수가 늘어난 것과 같습니다). 활성 레인 수를 절반으로 줄이면 활성 레인당 행 수가 두 배가 됩니다(예: 활성 4개 → 뱅크당 256개 행, 활성 1개 → 1024개 행).

더블 버퍼링

TRF는 각 뱅크를 두 절반으로 나누어 더블 버퍼링을 가능하게 합니다. TRF Sequencer가 한쪽 절반에서 로드하는 동안 저장이 다른 쪽 절반을 채우며, 반복 사이에 두 절반을 서로 바꿀 수 있습니다. 세 가지 주소 모드가 영역을 지정하며, 스케줄러는 저장 시점에 그중 하나를 고릅니다. Full은 뱅크당 128개 행 전체를, FirstHalf는 0–63번 행을, SecondHalf는 64–127번 행을 사용합니다. 절반 모드에서는 Slice당 용량이 40 KB로 제한됩니다.

main 컨텍스트와 sub 컨텍스트에 걸쳐 이 두 절반을 사용하는 튜닝 패턴은 스케줄링과 튜닝: 더블 버퍼링을 참조하십시오.

두 절반은 같은 뱅크를 공유하므로, 읽기와 쓰기는 서로 다른 행을 대상으로 하더라도 뱅크 수준에서 경합합니다. 한 사이클에 둘이 같은 뱅크를 대상으로 하면 읽기가 우선합니다. 축약 파이프라인은 이번 사이클에 데이터가 필요하지만 저장은 기다릴 수 있기 때문입니다.

TRF는 읽기 캐시와 뱅크 교대로 이 경합을 완화합니다.

TRF 읽기는 재사용이 많습니다. 같은 데이터가 보통 여러 사이클에 걸쳐 브로드캐스트되므로, 직접 매핑 읽기 캐시(8 레인 × 2 뱅크 × 4 행 × 320비트 = 2.5 KB)가 뱅크 앞에 위치하여 반복되는 읽기를 흡수합니다. 캐시는 동시에 수행되는 저장과의 경합도 줄여 줍니다. 히트 시에는 읽기가 뱅크를 거치지 않으므로, 저장이 그 사이클에 뱅크를 사용할 수 있습니다. 미스 시에는 캐시가 뱅크에서 다시 채워지며, 그 사이클 동안 뱅크를 점유합니다.

좁은 읽기(≤ 32바이트)의 경우, 뱅크 교대가 두 번째 완화 수단을 더합니다. 읽기는 뱅크 하나만 사용하므로, 두 뱅크에 걸쳐 32바이트 단위로 교대할 수 있습니다. 그러면 연속된 사이클에서 읽기와 쓰기가 서로 다른 뱅크에 놓이게 되어, 캐시 미스가 나더라도 경합을 피합니다. 넓은 읽기(64바이트)는 매 사이클 두 뱅크를 모두 점유하므로 캐시 미스가 날 때마다 동시에 수행되는 저장이 막히지만, 좁은 읽기는 캐시 미스가 나더라도 절반 대역폭의 교대를 유지합니다.

Vector Register File

VRF는 Collect Engine에서 쓰이거나 Data Memory에서 직접 쓰이며, Vector Engine이 읽습니다.

인터페이스

VrfTensor는 VRF에 저장된 텐서입니다:

/// Tensor stored in the vector register file (VRF).
#[primitive(VrfTensor)]
#[derive(Debug, Clone)]
pub struct VrfTensor<D: VeScalar, Chip: M, Cluster: M, Slice: M, Element: M, B: Backend = CurrentBackend> {
    pub(crate) inner: Tensor<D, Pair<Chip, Pair<Cluster, Pair<Slice, Element>>>, B>,
    _marker: PhantomData<(D, Chip, Cluster, Slice, Element)>,
}

Chip / Cluster / Slice는 소스에서 그대로 전달됩니다. Element는 (Slice별) 레이아웃을 담습니다.

Vector Engine 연산은 이 세 가지를 기준으로 피연산자를 대조하므로, 피연산자가 공급 대상 스트림과 다른 분할로 전달되면 다른 Slice의 데이터를 읽는 대신 컴파일 오류가 발생합니다(Operands 참조). .reshape::<Chip2, Cluster2, Slice2, Element2>()는 같은 Slice들의 다른 분할로 피연산자를 다시 기술하며, 익명 브로드캐스트(m![256])로 복제된 피연산자가 이름 있는 축으로 분할된 스트림에 공급되는 방식이 바로 이것입니다. 이 메서드는 매핑의 레이블만 바꾸고 데이터는 옮기지 않으므로, DmTensor::reshape와 같은 이유로 unsafe입니다. 즉, 모든 물리적 위치가 새 매핑이 주장하는 값을 이미 담고 있어야 합니다. 축을 재그룹화하는 것과 브로드캐스트 분배 축에 이름을 붙이는 것이 이를 만족하는 두 가지 형태이며, VrfTensor::reshape의 # Safety 섹션에서 둘 다 설명합니다.

Collect Engine에서

CollectTensor에 대해 .to_vrf::<Element2>()를 호출하면 플릿을 VRF에 저장하고 VrfTensor를 생성합니다. 플릿이 Register File의 어디에 놓일지는 컴파일러가 결정하므로, 호출에서 주소를 지정하지 않습니다:

impl<'l, P: CanApplyToVrf, D: VeScalar, Chip: M, Cluster: M, Slice: M, Time: M, Packet: M, B: Backend>
    TuTensor<'l, { Tu::Sub }, P, D, Chip, Cluster, Slice, Time, Packet, B>
{
    /// Stores to the vector register file.
    #[primitive(TuTensor::to_vrf)]
    pub fn to_vrf<Element: M>(self) -> VrfTensor<D, Chip, Cluster, Slice, Element, B> {
        verify_to_vrf::<D, Time, Packet, Element>();
        VrfTensor::new(self.inner.transpose(false))
    }
}

impl<'l, P: CanApplyToVrf, D: VeScalar, Chip: M, Cluster: M, Slice: M, Time: M, Packet: M, B: Backend>
    TuTensor<'l, { Tu::Main }, P, D, Chip, Cluster, Slice, Time, Packet, B>
{
    /// Stores to the vector register file, occupying the sub context the write runs through.
    ///
    /// The borrow lasts the statement. How long the *register* stays occupied is the compiler's
    /// allocation to make, so the returned [`VrfTensor`] keeps no borrow.
    #[primitive(TuTensor::to_vrf_from_main)]
    pub fn to_vrf<Element: M>(
        self,
        _sub: &mut TuContext<{ Tu::Sub }>,
    ) -> VrfTensor<D, Chip, Cluster, Slice, Element, B> {
        verify_to_vrf::<D, Time, Packet, Element>();
        VrfTensor::new(self.inner.transpose(false))
    }
}

.to_vrf는 스트리밍되는 Time / Packet을 Element2로 평탄화합니다:

Element2 = [Time, Packet]

Element2는 사용자가 선택합니다.

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

fn store_vrf<'l>(
    input: CollectTensor<'l, { Tu::Sub }, i32, m![1], m![1 # 2], m![1 # 256], m![B / 8], m![B % 8]>,
) -> VrfTensor<i32, m![1], m![1 # 2], m![1 # 256], m![B]> {
    input.to_vrf()
}

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

let c: CollectTensor<'_, { Tu::Sub }, i32, m![1], m![1 # 2], m![1 # 256], m![B / 8], m![B % 8]> = CollectTensor::new(&mut device.sub, Tensor::zero());
let _o = store_vrf(c);
}

저장은 컨텍스트마다 하나의 시그니처를 가지며, 위 예제는 sub 컨텍스트의 시그니처를 사용합니다. 메인 컨텍스트에서 발행한 저장은 device.sub를 인자로 받습니다. 쓰기 레지스터가 sub 컨텍스트의 레지스터 맵에 있기 때문입니다. 저장이 어느 컨텍스트를 점유하는가를 참고하십시오.

Data Memory에서

입력 접근이 완전히 연속적인 경우(간격이나 재정렬이 없는 경우), VRF는 짧은 명령(StoVRF)을 지원합니다. 이는 Fetch → Switch → Collect → to_vrf()로 이어지는 전체 파이프라인을 우회하여 Data Memory에서 VRF로 데이터를 직접 로드하는 간결한 하드웨어 명령어입니다. 이 지름길은 임의 레이아웃 지원을 포기하는 대신 설정 오버헤드를 줄입니다.

아키텍처