Packet Reducer

Packet Reducer는 lane마다 리덕션 트리 하나를 두고, 단일 Packet 안에서 가장 안쪽의 축약 축들을 reduce-add합니다.

인터페이스

.contract_packet()는 Packet Reducer를 호출합니다. 각 lane은 Outer 단계의 OutPacket에서 물려받은 i4, i8, f8 또는 bf16 요소로 이루어진 32 B 또는 64 B Packet을 받습니다. 형식적으로는 \(\text{output}[i] = \sum_{j} \text{input}[i, j]\)를 계산하며, 여기서 i는 남는 (출력) 축들에 걸쳐 있고 j는 Packet 내부의 축약 축들에 걸쳐 있습니다.

impl<
    'l,
    const T: Tu,
    D: Scalar,
    Storage: ContractionCast<Output = D>,
    Chip: M,
    Cluster: M,
    Slice: M,
    Lane: M,
    Time: M,
    Packet: M,
    B: Backend,
> ContractOuterTensor<'l, T, D, Storage, Chip, Cluster, Slice, Lane, Time, Packet, B>
{
    /// Spatial reduction within `Packet`: validates the reduce-add along the contracted axes inside
    /// `Packet` that the fused fold at `contract_lane` will perform. `D` is the widened accumulator the
    /// deferred carrier stays keyed on; the DPE input packet is still sized in `Storage` bytes.
    #[primitive(ContractOuterTensor::contract_packet)]
    pub fn contract_packet<OutPacket: M>(
        self,
    ) -> ContractPacketTensor<'l, T, D, Chip, Cluster, Slice, Lane, Time, OutPacket, B> {
        verify_contract_packet(ContractPacketInput {
            in_packet: Packet::to_value(),
            out_packet: OutPacket::to_value(),
            element_bits: Storage::BITS,
        });
        // Carry the deferred operands forward unreduced: the fused contraction at `contract_lane`
        // performs this Packet reduction too. This stage only re-types the carrier to `OutPacket`.
        ContractPacketTensor::new(self.device, self.inner)
    }
}

아래 커널은 8개 lane을 모두 병렬로 사용합니다. 트리 깊이 5로 B의 bf16 요소 32개를 합산하여 A 위치마다 f32 하나를 생성합니다.

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

fn matmul<'l, const T: Tu>(
    input: CollectTensor<'l, T, bf16, m![1], m![1 # 2], m![1 # 256], m![A, B / 16], m![B % 16]>,
    trf: &TrfTensor<bf16, m![1], m![1 # 2], m![1 # 256], m![C], m![B]>,
) -> ContractTensor<'l, T, f32, m![1], m![1 # 2], m![1 # 256], m![A], m![C]> {
    //
    // Spatial reduction: tree depth 5 reduces 32 bf16 elements along B → f32
    // Output (Interleaved): Time = [A], Packet = [C]
    input.contract_outer::<m![A], m![B], _, _, _>(&trf)
         .contract_packet::<m![1]>()
         .contract_time::<m![A]>()
         .contract_lane::<m![A], m![C]>(LaneMode::Interleaved)
}

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

let a: CollectTensor<'_, _, bf16, m![1], m![1 # 2], m![1 # 256], m![A, B / 16], m![B % 16]> = CollectTensor::new(&mut device.main, Tensor::zero());
let b: TrfTensor<bf16, m![1], m![1 # 2], m![1 # 256], m![C], m![B]> = TrfTensor::zero();
let _o = matmul(a, &b);
}

아키텍처

ReducePacket = Packet / 2^d        for 0 ≤ d ≤ log2(Packet::SIZE)
OutPacket    = ReducePacket, outermost padding removed
               (size ≤ 32; when ReducePacket::SIZE > 32 only the innermost 32 survive)

Packet Reducer는 먼저 입력 Packet에 대해 lane마다 독립적인 리덕션 트리를 실행합니다. 깊이 0에서 트리의 리프는 입력 Packet의 요소를 담고 있으며, 이후 각 깊이에서는 쌍을 합산하여 요소 수를 절반으로 줄입니다. 최대 트리 깊이는 log2(Packet::SIZE)이므로 i4(Packet::SIZE = 128)는 7, i8 / f8(64)는 6, bf16(32)는 5입니다. 사용자의 OutPacket이 주어지면 컴파일러가 트리 깊이 d를 도출하고, 트리는 가장 안쪽의 2^d개 요소를 소비하여 ReducePacket을 생성합니다.

그다음 Packet Reducer는 ReducePacket을 OutPacket으로 잘라내는데, 하류에 있는 Time Reducer의 lane별 누산기에는 열이 32개뿐이므로 OutPacket의 크기는 최대 32개 요소로 제한됩니다. ReducePacket::SIZE > 32이면 바깥쪽 더미가 잘려 나가고 가장 안쪽의 32개 요소만 남습니다. 예를 들어 i4는 128개 요소의 Packet으로 도착하므로 d ∈ {0, 1}은 128개 또는 64개 요소의 ReducePacket을 생성하며, 둘 다 OutPacket::SIZE = 32로 잘립니다.

ReducePacket은 가장 바깥쪽 패딩(패딩된 packet 축)을 가질 수도 있으며, Packet Reducer는 그 패딩을 버리고 유효한 열만 남깁니다. 예를 들어 [A # 16]의 ReducePacket(A는 유효하며 16개 열까지 패딩됨)은 OutPacket = [A]로 리덕션됩니다. 따라서 OutPacket은 패딩이 있는 상태([A # 16])로 선언할 수도 있고 이미 제거된 상태([A])로 선언할 수도 있으며, 둘 다 같은 결과를 가리킵니다.

성능

지연 시간은 트리 깊이에 따라 달라지며, i4는 7사이클, i8/f8는 6사이클, bf16는 5사이클입니다. 요소 타입이 넓을수록 각 Packet에 들어가는 요소 수가 적으므로 깊이가 줄어듭니다. 가산기 트리는 완전히 파이프라인화되어 있으므로 깊이는 첫 출력까지의 지연 시간을 늘리지만 정상 상태의 처리량을 떨어뜨리지는 않습니다. 파이프라인이 채워지고 나면 매 사이클마다 Packet 하나가 들어가고 리덕션된 출력 하나가 나옵니다.

Lane < 8이면 비활성 lane의 리덕션 트리는 유휴 상태가 되므로 사이클당 처리량은 Lane::SIZE / 8에 비례하여 감소합니다.