Inter-Slice Reducer

Inter-Slice Reducer는 Cluster 안의 256개 Slice에 걸쳐 텐서를 리덕션합니다. Chip, Cluster, Packet은 그대로 유지하고, Slice와 Time은 OutSlice와 OutTime으로 바꿔 씁니다. 출력 텐서는 입력 모드와 관계없이 항상 Way8입니다.

인터페이스

리듀서에는 vector_init() 직후에 진입하거나(아래에 나오는 InterFirst 경로), 호환되는 intra-slice 단계에서 진입할 수 있습니다(intra-slice 텐서에 같은 vector_inter_slice_reduce() 메서드를 호출하는 IntraFirst 경로). 아래에 나오는 시그니처는 VectorInitTensor 변형입니다. 이 전환을 지원하는 단계의 intra-slice 텐서에도 같은 메서드가 있으므로 호출 코드의 모양은 똑같습니다.

Inter-Slice Reducer는 i32와 f32에 대해 별도의 API를 제공합니다.

`i32` 연산

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)
    }
}

InterSliceReduceOpI32 연산:

연산	설명
`Add`	래핑 덧셈
`AddSat`	포화 덧셈
`Max`	최댓값
`Min`	최솟값

`f32` 연산

impl<'l, const T: Tu, Chip: M, Cluster: M, Slice: M, Time: M, Packet: M>
    VectorInitTensor<'l, T, f32, Chip, Cluster, Slice, Time, Packet>
{
    /// Performs inter-slice reduce for f32 as the first VE operation.
    #[primitive(VectorInitTensor::vector_inter_slice_reduce)]
    pub fn vector_inter_slice_reduce<OutSlice: M, OutTime: M>(
        self,
        op: InterSliceReduceOpF32,
    ) -> VectorInterSliceReduceTensor<'l, T, f32, 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)
    }
}

InterSliceReduceOpF32 연산:

연산	설명
`Add`	부동소수점 덧셈
`Max`	최댓값
`Min`	최솟값
`Mul`	부동소수점 곱셈

제약 조건

지원되는 Slice → OutSlice 및 Time → OutTime 형태는 다음 네 가지 규칙을 따릅니다.

가장 안쪽부터 리덕션하십시오. Slice에서 리덕션되는 부분은 가장 안쪽 인수들이어야 하며, 연속적이어야 하고, 리덕션 비율 r까지 스트라이드가 1이어야 합니다.
리덕션된 축의 대체. OutSlice에서 리덕션된 각 인자의 자리는 더미(1 # n), 새 차원에 대한 브로드캐스트, Time으로부터의 승격 중 하나로 채워집니다.
대체 종류는 자유롭게 섞을 수 있습니다. 더미, 브로드캐스트, 승격 자리는 OutSlice에 순서와 관계없이 함께 나타날 수 있습니다.
Time에서 OutSlice로의 승격은 순서를 바꿉니다. Time → OutTime 부분은 남아 있는 인자들의 상대적 순서를 유지하지만, Time → OutSlice 승격 경로는 순서를 유지하지 않습니다. 승격된 인자의 OutSlice 내 위치는 Time 내 위치와 무관합니다.

예제

아래 예제의 수식은 einsum 표기법을 사용합니다. 입력 쪽에는 나타나지만 출력 쪽에는 나타나지 않는 차원은 리덕션(합산)되고, 출력 쪽에는 나타나지만 입력 쪽에는 나타나지 않는 차원은 브로드캐스트됩니다.

더미 대체

이 패스는 입력을 R에 걸쳐 합산하고 그 결과를 더미 자리에 배치합니다. einsum 형식은 AR -> A입니다.

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

// When R is reduced and no other dimension fills its slot, the output keeps the slot as a 1 # n dummy.
// One position holds the reduced value, and the remaining n - 1 are padding.
// `# n` denotes dimension multiplicity (see the Mapping Expressions doc).
fn inter_slice_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, the freed R-slot becomes the 1 # 4 dummy in OutSlice
        .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 = inter_slice_add(c);
}

Slice = [A / 8, R]  ->  [A / 8, 1 # 4]
Time  = [1]         ->  [1]

새 Slice 차원으로의 브로드캐스트

이 패스는 R을 리덕션하고 그 결과를 새 차원 X에 대해 브로드캐스트합니다. einsum 형식은 PRW -> PWX이며, 출력 쪽의 새 X가 브로드캐스트를 나타냅니다.

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

// A fresh non-reduce dimension X takes the slot that R leaves behind.
// The reduced value broadcasts across every position of X.
fn broadcast_into_x<'l, const T: Tu>(
    input: CollectTensor<'l, T, f32, m![1], m![B], m![W, R], m![1], m![P]>,
) -> VectorFinalTensor<'l, T, f32, m![1], m![B], m![W, X], m![1], m![P]> {
    input
        .vector_init()
        // sum across R, broadcast result over X (fresh OutSlice dimension)
        .vector_inter_slice_reduce::<m![W, X], m![1]>(InterSliceReduceOpF32::Add)
        .vector_final()
}

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

let c: CollectTensor<'_, _, f32, m![1], m![B], m![W, R], m![1], m![P]> = CollectTensor::new(&mut device.main, Tensor::zero());
let _o = broadcast_into_x(c);
}

Slice = [W, R]  ->  [W, X]
Time  = [1]     ->  [1]

`Time`에서 `OutSlice`로의 승격

이 패스는 R을 리덕션하고 Time 차원 V를 OutSlice로 승격합니다. einsum 형식은 PRSUVW -> PSUVW입니다.

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![B = 2, P = 8, R = 4, S = 2, U = 2, V = 4, W = 64];

// A dimension from Time (here V) is promoted into OutSlice to fill R's slot.
// The promoted dimension does not need to be outermost in Time.
fn axis_promotion<'l, const T: Tu>(
    input: CollectTensor<'l, T, f32, m![1], m![B], m![W, R], m![S, V, U], m![P]>,
) -> VectorFinalTensor<'l, T, f32, m![1], m![B], m![W, V], m![S, U], m![P]> {
    input
        .vector_init()
        // sum across R, V moves from Time to OutSlice
        .vector_inter_slice_reduce::<m![W, V], m![S, U]>(InterSliceReduceOpF32::Add)
        .vector_final()
}

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

let c: CollectTensor<'_, _, f32, m![1], m![B], m![W, R], m![S, V, U], m![P]> = CollectTensor::new(&mut device.main, Tensor::zero());
let _o = axis_promotion(c);
}

Slice = [W, R]     ->  [W, V]
Time  = [S, V, U]  ->  [S, U]

리듀서와 Intra-Slice Chain을 어느 순서로든 결합한 예제는 Vector Engine 예제를 참고하십시오.

성능

리덕션 비율 r(하나의 리덕션 그룹에 속한 Slice 수)이 주요 튜닝 요소입니다. Inter-Slice 리덕션 지연 시간은 O(r) 사이클로, 리덕션 그룹을 링으로 대략 한 바퀴 순회하는 시간입니다. 총 시간은 입력 스트리밍 시간에 링 크기만큼의 꼬리 구간을 더한 값입니다.

실제로는 업스트림 작업(부분합을 생성하는 축약 또는 vector_inter_slice_reduce() 이전의 intra-slice 작업)이 대개 더 오래 걸려 이 링 꼬리 구간을 가리므로, 리듀서가 병목이 되지 않습니다. 리듀서의 영향은 r이 클 때(꼬리 구간이 더 길어짐)와 고정된 꼬리 구간을 여러 패킷에 걸쳐 분산할 수 없는 작은 텐서에서 드러납니다.

Keyboard shortcuts

Tensor Contraction Processor 프로그래밍