Intra-Slice Reduce

Intra-Slice Chain의 IntraSliceReduce 단계는 각 Slice의 Time과 Packet에 있는 차원을 리덕션합니다(Chip, Cluster, Slice는 변경 없이 그대로 통과합니다). 이와 상보적인 경우인 Cluster의 256개 Slice에 걸친 리덕션은 Inter-Slice Reducer가 담당합니다.

예제

reduce 호출의 주요 매개변수는 다음과 같습니다.

REDUCE_LABEL: 리덕션할 축입니다. 리덕션은 이 축을 담고 있는 Time과 Packet의 모든 인자를 제거하므로, 이 인자들은 출력 형태(OutTime, OutPacket)에 나타나서는 안 됩니다. 예를 들어 R이 Time에서 R / 4로, Packet에서 R % 4로 분할된 경우, REDUCE_LABEL = R을 지정하면 둘 다 제거됩니다.
op: 리덕션 연산입니다. IntraSliceReduceOpI32는 AddSat, Max, Min을 제공하고, IntraSliceReduceOpF32는 Add, Max, Min을 제공합니다.
OutTime, OutPacket: 리덕션 후의 출력 Time 및 Packet 형태입니다. 이는 입력 Time과 Packet에서 모든 REDUCE_LABEL 인자를 제거한 것과 같습니다.

이 호출은 Pair Mode에서는 사용할 수 없습니다. 아래 예제는 각 매개변수 조합을 다룹니다.

`Time`에서의 리덕션

R는 Time에만 존재하므로 이 단계는 시간 단계에 걸쳐 누산합니다. 포화 덧셈으로 \(output[a] = \sum_{r \in R} input[a, r]\)를 계산합니다.

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

// R in Time → temporal accumulation. Packet is non-reduce.
fn reduce_time<'l, const T: Tu>(
    input: VectorBranchTensor<'l, T, i32, m![1], m![1 # 2], m![A / 2], m![R], m![A % 2 # 8], Fresh, { stage::VeOrder::IntraFirst }>,
) -> VectorIntraSliceReduceTensor<'l, T, i32, m![1], m![1 # 2], m![A / 2], m![1], m![A % 2 # 4], Fresh, { stage::VeOrder::IntraFirst }>
{
    input
        .vector_narrow_trim::<m![A % 2 # 4]>()       // 8-way → 4-way
        // R eliminated from Time
        .vector_intra_slice_reduce::<R, m![1], m![A % 2 # 4]>(
            IntraSliceReduceOpI32::AddSat,
        )
}

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

let i: VectorBranchTensor<'_, _, i32, m![1], m![1 # 2], m![A / 2], m![R], m![A % 2 # 8], Fresh, { stage::VeOrder::IntraFirst }> = VectorBranchTensor::new(&mut device.main, Tensor::zero(), TagMode::Zero);
let _o = reduce_time(i);
}

`Packet`에서의 리덕션

R이 Packet에만 존재하므로, 하드웨어는 각 플릿 내에서 4-way 트리 리덕션을 수행하고 시간 방향 누적은 건너뜁니다. \(output[a] = \sum_{r \in R} input[a, r]\)를 계산합니다.

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

// R in Packet → tree reduce within flit.
fn reduce_packet<'l, const T: Tu>(
    input: VectorBranchTensor<'l, T, f32, m![1], m![1 # 2], m![A / 2], m![A % 2], m![R # 8], Fresh, { stage::VeOrder::IntraFirst }>,
) -> VectorIntraSliceReduceTensor<'l, T, f32, m![1], m![1 # 2], m![A / 2], m![A % 2], m![1 # 4], Fresh, { stage::VeOrder::IntraFirst }>
{
    input
        .vector_narrow_trim::<m![R]>()             // 8-way → 4-way
        // R eliminated from Packet
        .vector_intra_slice_reduce::<R, m![A % 2], m![1 # 4]>(
            IntraSliceReduceOpF32::Add,
        )
}

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

let i: VectorBranchTensor<'_, _, f32, m![1], m![1 # 2], m![A / 2], m![A % 2], m![R # 8], Fresh, { stage::VeOrder::IntraFirst }> = VectorBranchTensor::new(&mut device.main, Tensor::zero(), TagMode::Zero);
let _o = reduce_packet(i);
}

양쪽 모두에서의 리덕션

R는 Packet과 Time에 나뉘어 분할됩니다. Packet의 R % 4는 각 플릿 내에서 트리 리덕션을 수행하고, Time의 R / 4는 시간 단계에 걸쳐 누산합니다. \(output[a] = \max_{r \in R} input[a, r]\)를 계산합니다.

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

// R % 4 in Packet → spatial tree reduce
// R / 4 in Time → temporal accumulation
fn reduce_time_packet<'l, const T: Tu>(
    input: VectorBranchTensor<'l, T, f32, m![1], m![1 # 2], m![A], m![R / 4], m![R % 4 # 8], Fresh, { stage::VeOrder::IntraFirst }>,
) -> VectorIntraSliceReduceTensor<'l, T, f32, m![1], m![1 # 2], m![A], m![1], m![1 # 4], Fresh, { stage::VeOrder::IntraFirst }>
{
    input
        .vector_narrow_trim::<m![R % 4]>()            // 8-way → 4-way
        // R eliminated from both Time and Packet
        .vector_intra_slice_reduce::<R, m![1], m![1 # 4]>(
            IntraSliceReduceOpF32::Max,
        )
}

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

let i: VectorBranchTensor<'_, _, f32, m![1], m![1 # 2], m![A], m![R / 4], m![R % 4 # 8], Fresh, { stage::VeOrder::IntraFirst }> = VectorBranchTensor::new(&mut device.main, Tensor::zero(), TagMode::Zero);
let _o = reduce_time_packet(i);
}

R은 Slice, Time, Packet에 나뉘어 있지만, Slice 내부 리듀서는 Time과 Packet 부분만 축소하므로 R의 Slice 부분은 출력에 남습니다. 여기서 R = 13은 레이아웃에 맞추기 위해 32로 패딩되고(R # 32), 이후 R은 4개의 Slice에 걸쳐 분할됩니다(Slice당 R 위치 8개). 위치 0-12만 실제 요소를 담고 있으므로, 그 경계에 걸친 Slice(위치 8-15: 실제 요소 5개 뒤에 패딩 3개)는 경계 Slice이며, 그 이후의 Slice는 전부 패딩입니다. VCG가 Slice별 리덕션 횟수를 제어하여 각 Slice가 자신의 실제 요소만 리덕션하도록 합니다(정확한 매핑은 Valid Count Generator를 참조하십시오).

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

// R split across all three: Slice (groups of 8), Time (pairs within group), Packet (4 elements).
fn reduce_slice_time_packet<'l, const T: Tu>(
    input: VectorBranchTensor<'l, T, i32, m![1], m![1 # 2], m![R # 32 / 8 # 256], m![R # 32 / 4 % 2], m![R # 32 % 4 # 8], Fresh, { stage::VeOrder::IntraFirst }>,
) -> VectorIntraSliceReduceTensor<'l, T, i32, m![1], m![1 # 2], m![R # 32 / 8 # 256], m![1], m![1 # 4], Fresh, { stage::VeOrder::IntraFirst }>
{
    input
        .vector_narrow_trim::<m![R # 32 % 4]>()       // 8-way → 4-way
        // R eliminated from Time and Packet (accumulated within each slice)
        .vector_intra_slice_reduce::<R, m![1], m![1 # 4]>(
            IntraSliceReduceOpI32::Min,
        )
}

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

let i: VectorBranchTensor<'_, _, i32, m![1], m![1 # 2], m![R # 32 / 8 # 256], m![R # 32 / 4 % 2], m![R # 32 % 4 # 8], Fresh, { stage::VeOrder::IntraFirst }> = VectorBranchTensor::new(&mut device.main, Tensor::zero(), TagMode::Zero);
let _o = reduce_slice_time_packet(i);
}

아키텍처

이 단계는 Time 축과 Packet 축에 대해 별도의 장치를 운용합니다.

`Time`에서의 리덕션

이 단계는 슬롯 용량이 8인 시간 누산기 모델을 적용합니다(따라서 InnerTime::SIZE ≤ 8입니다).

위의 reduce_time에서는 Time = m![R]이고 OutTime = m![1]이므로 R이 가장 바깥쪽 리덕션 차원이며 InnerTime = m![1](InnerTime::SIZE = 1)입니다. 단일 슬롯이 모든 R 값을 출력에 누산합니다.

InnerTime::SIZE가 8을 초과하면 API가 호출을 거부합니다. 예를 들면 다음과 같습니다.

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 6, B = 16, R = 16];
fn invalid_too_many_slots<'l, const T: Tu>(
    input: VectorBranchTensor<'l, T, i32, m![1], m![1 # 2], m![A / 3 # 256], m![R, A % 3, B % 4], m![B / 4 # 8], Fresh, { stage::VeOrder::IntraFirst }>,
) -> VectorIntraSliceReduceTensor<'l, T, i32, m![1], m![1 # 2], m![A / 3 # 256], m![A % 3, B % 4], m![B / 4], Fresh, { stage::VeOrder::IntraFirst }>
{
    input
        .vector_narrow_trim::<m![B / 4]>()
        // Time      = m![R, A % 3, B % 4]
        // OutTime   = m![A % 3, B % 4]
        // InnerTime = m![A % 3, B % 4], InnerTime::SIZE = 3 × 4 = 12 > 8
        .vector_intra_slice_reduce::<R, m![A % 3, B % 4], m![B / 4]>(
            IntraSliceReduceOpI32::AddSat,
        )
    // Rejected: 12 accumulator slots required, but only 8 are available.
}

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

let i: VectorBranchTensor<'_, _, i32, m![1], m![1 # 2], m![A / 3 # 256], m![R, A % 3, B % 4], m![B / 4 # 8], Fresh, { stage::VeOrder::IntraFirst }> = VectorBranchTensor::new(&mut device.main, Tensor::zero(), TagMode::Zero);
let _o = invalid_too_many_slots(i);
}

`Packet`에서의 리덕션

플릿당 4개의 Packet 요소는 다음 두 경로 중 하나를 따릅니다.

OutPacket = Packet: 요소가 리덕션되지 않고 그대로 통과합니다. 사이클당 4개가 출력되며, 각 출력은 Time에 걸쳐 독립적으로 누산됩니다.
OutPacket = m![1 # 4]: 요소가 2단계 트리 op(op(a, b), op(c, d))를 거쳐 단일 값으로 축약됩니다. 사이클당 1개가 출력되며 패딩 위치 3개가 추가됩니다.

리덕션 축을 하드웨어 차원에 맞추기 위해 패딩하면, 패딩된 위치에는 리덕션에서 제외해야 하는 임의의 데이터가 들어 있습니다. 패딩을 제외하는 전략은 두 가지입니다.

VCG (Valid Count Generator): 해당 축 배치가 지원되는 경우 이 방식을 우선적으로 사용합니다. 컴파일러가 매핑으로부터 VCG를 자동으로 구성하고, VCG는 각 플릿에 valid_count를 태그로 붙여 패딩 요소가 자동으로 제외되도록 합니다. Slice, Time, Packet에 걸친 모든 축 배치가 지원되는 것은 아닙니다. 자세한 내용은 Valid Count Generator를 참고하십시오.
항등 요소 패딩: 데이터가 Intra-Slice Chain에 도달하기 전에 패딩 위치를 리덕션 연산의 항등 요소로 채우십시오. 이 책에는 사용할 수 있는 공개 마스킹 연산이 문서화되어 있지 않으므로, 리덕션 전에 항등원을 생성하는 검증된 경로를 사용하십시오.

연산 항등 요소

AddSat / Add 0 / 0.0

Max i32::MIN / f32::NEG_INFINITY

Min i32::MAX / f32::INFINITY

이 전략은 리덕션 연산 앞에 역변환이 불가능한 변환이 없는 경우에만 적용됩니다. 예를 들어 exp(x) + exp(y) + ...(지수 함수의 합)에서는 exp(p) = 0(덧셈 항등원)을 만족하는 값 p가 존재하지 않습니다. 따라서 항등원 패딩을 적용할 수 없습니다.

연산	항등 요소
`AddSat` / `Add`	`0` / `0.0`
`Max`	`i32::MIN` / `f32::NEG_INFINITY`
`Min`	`i32::MAX` / `f32::INFINITY`

성능

트리 리덕션이 Intra-Slice Chain 안에서 완전히 파이프라인화되어 플릿당 추가 비용이 없으므로, 처리량은 사이클당 플릿 하나로 유지됩니다.

이 단계는 결과를 내보내기 전에 리덕션 그룹의 모든 입력 플릿을 누산해야 하므로, 지연 시간에는 n 플릿 사이클만큼의 첫 출력 지연이 추가됩니다. 여기서 n은 리덕션 축의 시간 단계 수입니다. 여러 엔진으로 구성된 파이프라인에서는 이 누산 지연 때문에 첫 플릿을 기다리는 다운스트림 엔진이 정체됩니다.

Keyboard shortcuts

Tensor Contraction Processor 프로그래밍