Vector Engine

Vector Engine은 요소별 연산과 리덕션을 수행합니다. 활성화 함수(GELU, SiLU), 정규화(softmax, layer norm), 이항 연산, Slice 내부 및 Slice 간 리덕션 등이 그 예입니다.

이 엔진은 32비트 타입인 i32와 f32만 받습니다. 앞 단계의 Contraction Engine은 타입을 자동으로 확장합니다(bf16 곱은 f32로, i8 곱은 i32로 누적됩니다). 이 엔진을 거치지 않는 경우에는 Fetch Adapter가 타입 캐스트 어댑터로 입력을 확장해야 합니다.

인터페이스

한 번의 Tensor Unit 호출에서 Vector Engine 부분은 vector_init()부터 vector_final()까지의 메서드 체인입니다. 이 엔진은 두 하위 구성 요소로 이루어집니다. 하나는 Intra-Slice Chain(요소별 / 이항 / Slice별 리덕션 단계)이고, 다른 하나는 Inter-Slice Reducer(Cluster 내 256개 Slice에 걸쳐 리덕션)입니다. vector_init()와 vector_final() 사이에서는 체인만 실행하거나, 리듀서만 실행하거나, 둘 다 실행할 수 있습니다. 둘 다 실행하는 경우 순서는 IntraFirst(체인 다음 리듀서) 또는 InterFirst(리듀서 다음 체인)입니다.

Intra-Slice Chain은 vector_intra_slice_tag()로 진입하며, 입력에 두 개의 병렬 스트림으로 나눌 2-way 그룹화 축이 있는 경우에는 vector_intra_slice_unzip()로 진입합니다(Pair Mode 참조). 두 진입점 모두 vector_init() 직후 또는 Inter-Slice Reducer의 출력에서 실행됩니다. Inter-Slice Reducer는 vector_init() 직후 또는 호환되는 Intra-Slice 단계에서 vector_inter_slice_reduce()로 진입합니다. 단계별 API 설명은 Intra-Slice Chain과 Inter-Slice Reducer를 참조하십시오.

아래 시그니처는 vector_init() 쪽 진입 메서드만 다룹니다. 같은 이름의 메서드(vector_intra_slice_tag, vector_inter_slice_reduce)는 체인↔리듀서 전환을 위해 체인 텐서와 리듀서 텐서에도 있으며, 이는 하위 페이지에서 설명합니다.

impl<'l, const T: Tu, P: CanApplyVectorInit, D: VeScalar, Chip: M, Cluster: M, Slice: M, Time: M, Packet: M>
    TuTensor<'l, T, P, D, Chip, Cluster, Slice, Time, Packet>
{
    /// Initializes Vector Engine processing for this tensor.
    #[primitive(TuTensor::vector_init)]
    pub fn vector_init(self) -> VectorInitTensor<'l, T, D, Chip, Cluster, Slice, Time, Packet> {
        VectorInitTensor::new(self.device, self.inner)
    }
}

impl<'l, const T: Tu, D: VeScalar, Chip: M, Cluster: M, Slice: M, Time: M, Packet: M>
    VectorInitTensor<'l, T, D, Chip, Cluster, Slice, Time, Packet>
{
    /// Enters VE intra-slice pipeline (single stream).
    #[primitive(VectorInitTensor::vector_intra_slice_tag)]
    pub fn vector_intra_slice_tag(
        self,
        branch: TagMode<D>,
    ) -> VectorBranchTensor<'l, T, D, Chip, Cluster, Slice, Time, Packet, Fresh, { VeOrder::IntraFirst }> {
        VectorBranchTensor::new(self.device, self.inner, branch)
    }

    /// Enters VE intra-slice pipeline (two-group / unzip).
    #[primitive(VectorInitTensor::vector_intra_slice_unzip)]
    pub fn vector_intra_slice_unzip<I: AxisName, SplitTime: M>(
        self,
    ) -> VectorTensorPair<'l, T, D, stage::Tag, Chip, Cluster, Slice, SplitTime, Packet> {
        verify_vector_intra_slice_unzip::<I, Time, SplitTime, Packet>();
        VectorTensorPair::new::<I, Time>(self.device, self.inner)
    }
}

impl<'l, const T: Tu, Chip: M, Cluster: M, Slice: M, Time: M, Packet: M>
    VectorInitTensor<'l, T, i32, Chip, Cluster, Slice, Time, Packet>
{
    /// Performs inter-slice reduce for i32 as the first VE operation.
    #[primitive(VectorInitTensor::vector_inter_slice_reduce)]
    pub fn vector_inter_slice_reduce<OutSlice: M, OutTime: M>(
        self,
        op: InterSliceReduceOpI32,
    ) -> VectorInterSliceReduceTensor<'l, T, i32, Chip, Cluster, OutSlice, OutTime, Packet, { VeOrder::InterFirst }>
    {
        let reduced = self.inner.reduce(op.reduce_fn(), op.identity(), true);
        create_inter_slice_reduce_tensor(self.device, reduced)
    }
}

예제

다음은 Vector Engine 호출이 어떤 모습인지 감을 잡을 수 있도록 보여 주는 몇 가지 대표 예제입니다. 전체 API 안내는 하위 페이지인 Intra-Slice Chain과 Inter-Slice Reducer에서 다룹니다.

ReLU 활성화

이 패스는 ReLU를 요소별로 적용하여 \(output[b, k, m, n] = \max(input[b, k, m, n], 0)\)를 계산합니다.

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![B = 2, K = 256, M = 16, N = 16];

// ReLU activation after batched matrix multiplication.
// Chain-only pass, so the reducer is skipped and the path trivially resolves to IntraFirst.
// Both clusters (B = 2) and all 256 slices (K) carry real data, no padding.
fn relu<'l, const T: Tu>(
    input: ContractTensor<'l, T, f32, m![1], m![B], m![K], m![M, N / 8], m![N % 8]>,
) -> VectorFinalTensor<'l, T, f32, m![1], m![B], m![K], m![M, N / 8], m![N % 8]> {
    input
        .vector_init()
        .vector_intra_slice_tag(TagMode::Zero)
        // max(x, 0), the ReLU itself
        .vector_clip(ClipBinaryOpF32::Max, 0.0f32)
        .vector_final()
}

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

let c: ContractTensor<'_, _, f32, m![1], m![B], m![K], m![M, N / 8], m![N % 8]> = ContractTensor::new(&mut device.main, Tensor::zero());
let _o = relu(c);
}

ReLU 후 리덕션

이 패스는 Slice별로 ReLU를 적용한 다음 R에 걸쳐 리덕션하여 \(output[a, b] = \sum_{r \in R} \max(input[a, b, r], 0)\)를 얻습니다.

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

// Chain applies ReLU, then reducer sums across slices.
// IntraFirst shape (chain runs first, then reducer).
// Both clusters (B = 2), all 256 slices (A / 8 * R), and full Way8 packet (A % 8) carry real data.
fn relu_then_reduce<'l, const T: Tu>(
    input: CollectTensor<'l, T, i32, m![1], m![B], m![A / 8, R], m![1], m![A % 8]>,
) -> VectorFinalTensor<'l, T, i32, m![1], m![B], m![A / 8, 1 # 4], m![1], m![A % 8]> {
    input
        .vector_init()
        .vector_intra_slice_tag(TagMode::Zero)
        // max(x, 0), the ReLU
        .vector_clip(ClipBinaryOpI32::Max, 0)
        // sum across R slices
        .vector_inter_slice_reduce::<m![A / 8, 1 # 4], m![1]>(InterSliceReduceOpI32::AddSat)
        .vector_final()
}

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

let c: CollectTensor<'_, _, i32, m![1], m![B], m![A / 8, R], m![1], m![A % 8]> = CollectTensor::new(&mut device.main, Tensor::zero());
let _o = relu_then_reduce(c);
}

리덕션 후 바이어스

이 패스는 R에 걸쳐 리덕션한 다음 상수 바이어스를 더하여 \(output[a, b] = \left(\sum_{r \in R} input[a, b, r]\right) + 100\)를 얻습니다.

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

// Reducer sums across slices, then chain adds a bias to the reduced result.
// InterFirst shape (reducer runs first, then chain).
// Both clusters (B = 2), all 256 slices, and full Way8 packet carry real data.
fn reduce_then_add<'l, const T: Tu>(
    input: CollectTensor<'l, T, i32, m![1], m![B], m![A / 8, R], m![1], m![A % 8]>,
) -> VectorFinalTensor<'l, T, i32, m![1], m![B], m![A / 8, 1 # 4], m![1], m![A % 8]> {
    input
        .vector_init()
        // sum across R slices
        .vector_inter_slice_reduce::<m![A / 8, 1 # 4], m![1]>(InterSliceReduceOpI32::AddSat)
        .vector_intra_slice_tag(TagMode::Zero)
        // add bias 100
        .vector_fxp(FxpBinaryOp::AddFxp, 100)
        .vector_final()
}

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

let c: CollectTensor<'_, _, i32, m![1], m![B], m![A / 8, R], m![1], m![A % 8]> = CollectTensor::new(&mut device.main, Tensor::zero());
let _o = reduce_then_add(c);
}

Slice 내부 및 Slice 간 리덕션

이 패스는 Slice 내부 리듀서(R의 Time 및 Packet 부분 담당)와 Inter-Slice Reducer(R의 Slice 부분 담당)를 결합하여 R을 완전히 리덕션합니다. einsum 형식은 BR -> B이며, 포화 덧셈을 사용합니다.

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

// R splits across Slice (R / 32 = 256), Time (R % 32 / 4 = 8), and Packet (R % 4, padded to 8 in Way8).
// Chain runs intra-slice reduce over R's Time and Packet portions, then the reducer collapses the Slice portion.
// IntraFirst shape (chain runs first, then reducer).
fn full_sum<'l, const T: Tu>(
    input: CollectTensor<'l, T, i32, m![1], m![B], m![R / 32], m![R % 32 / 4], m![R % 4 # 8]>,
) -> VectorFinalTensor<'l, T, i32, m![1], m![B], m![1 # 256], m![1], m![1 # 8]> {
    input
        .vector_init()
        .vector_intra_slice_tag(TagMode::Zero)
        // Way8 → Way4 (back 4 packet positions were padding)
        .vector_narrow_trim::<m![R % 4]>()
        // sum over R's Time and Packet portions
        .vector_intra_slice_reduce::<R, m![1], m![1 # 4]>(IntraSliceReduceOpI32::AddSat)
        // Way4 → Way8
        .vector_widen_pad::<m![1 # 8]>()
        // sum over R's Slice portion across all 256 slices
        .vector_inter_slice_reduce::<m![1 # 256], m![1]>(InterSliceReduceOpI32::AddSat)
        .vector_final()
}

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

let c: CollectTensor<'_, _, i32, m![1], m![B], m![R / 32], m![R % 32 / 4], m![R % 4 # 8]> = CollectTensor::new(&mut device.main, Tensor::zero());
let _o = full_sum(c);
}

쌍 덧셈

이 패스는 I를 따라 교차 배치된 두 그룹을 분리한 다음 쌍별로 더합니다. einsum 형식은 ABI -> AB입니다.

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

// Pair-mode entry via unzip, then a zip op fuses the two streams with an add.
// Both clusters (B = 2), all 256 slices (A / 8), and full Way8 packet (A % 8) carry real data.
fn pair_add<'l, const T: Tu>(
    input: CollectTensor<'l, T, i32, m![1], m![B], m![A / 8], m![I], m![A % 8]>,
) -> VectorFinalTensor<'l, T, i32, m![1], m![B], m![A / 8], m![1], m![A % 8]> {
    input
        .vector_init()
        // split into group 0 and group 1 along I
        .vector_intra_slice_unzip::<I, m![1]>()
        // group0 + group1
        .vector_clip_zip(ClipBinaryOpI32::AddFxp)
        .vector_final()
}

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

let c: CollectTensor<'_, _, i32, m![1], m![B], m![A / 8], m![I], m![A % 8]> = CollectTensor::new(&mut device.main, Tensor::zero());
let _o = pair_add(c);
}