Lane Folder

Lane Folder는 Contraction Engine의 마지막 단계입니다. Lane Folder는 Lane 차원의 8개 값을 OutPacket(Interleaved) 또는 OutTime(Sequential)으로 옮겨 이 차원을 제거합니다. 값을 합산하지는 않습니다. 이 단계는 Lane을 리덕션하지 않고 다른 축으로 접어 넣습니다.

인터페이스

.contract_lane(mode)는 Lane Folder를 호출합니다. 이 단계는 상류 Time Reducer의 버퍼를 8개 요소 폭의 출력 버스로 한 사이클씩 비우며, LaneMode는 각 사이클의 플릿에 무엇을 실을지 선택합니다.

impl<'l, const T: Tu, D: ContractionAccumulator, Chip: M, Cluster: M, Slice: M, Lane: M, Time: M, Packet: M, B: Backend>
    ContractTimeTensor<'l, T, D, Chip, Cluster, Slice, Lane, Time, Packet, B>
{
    /// Folds the `Lane` dimension into the output stream.
    /// `LaneMode::Interleaved` relocates `Lane` into `OutPacket`;
    /// `LaneMode::Sequential` relocates `Lane` into `OutTime`.
    #[primitive(ContractTimeTensor::contract_lane)]
    pub fn contract_lane<OutTime: M, OutPacket: M>(
        self,
        mode: LaneMode,
    ) -> ContractTensor<'l, T, D, Chip, Cluster, Slice, OutTime, OutPacket, B> {
        verify_contract_lane(ContractLaneInput {
            in_lane: Lane::to_value(),
            in_time: Time::to_value(),
            in_packet: Packet::to_value(),
            out_time: OutTime::to_value(),
            out_packet: OutPacket::to_value(),
            pre_reduce_time: self.pre_reduce_time,
            mode: mode.into(),
        });
        // Earlier stages only retype the deferred operands. Perform their fused contraction here,
        // then apply the lane-fold relayout.
        // `contraction_prewidened` validates `out` through mapping carve. Symbol-wise comparison is
        // invalid because one contracted symbol may span spatial and reduced slots.
        let contraction = self.inner;
        let out = <m![{ Chip }, { Cluster }, { Slice }, { Lane }, { Time }, { Packet }]>::to_value();
        let reduced: Tensor<D, m![{ Chip }, { Cluster }, { Slice }, { Lane }, { Time }, { Packet }], B> =
            Tensor::from_inner(B::contraction_prewidened(
                &contraction.lhs,
                &contraction.rhs,
                &contraction.lhs_map,
                &contraction.rhs_map,
                &contraction.pre_reduce,
                &out,
            ));
        ContractTensor::new(self.device, reduced.transpose(false))
    }
}

아래의 최소 예제는 ContractTimeTensor(상류 Time Reducer의 출력)를 받아 .contract_lane(...)만 호출하므로, 각 예제에서 Lane Folder만 따로 떼어 볼 수 있습니다. 입력 Packet의 크기는 Packet Reducer를 거친 뒤 남은 크기로, 레인당 {1, 2, 4, 8, 16, 32}개 요소 중 하나입니다.

Interleaved

Lane 차원은 OutPacket으로 접혀 들어갑니다. 각 사이클은 8개 레인 전체에서 하나의 열 위치를 읽으며(레인당 값 1개, 플릿당 값 8개), Lane은 가장 안쪽의 OutPacket으로 구체화됩니다.

OutTime   = [Time, Packet]
OutPacket = [Lane # 8]

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![N = 8, M = 4, P = 16];

/// Lane folds into OutPacket.
fn lane_interleaved<'l, const T: Tu>(
    // Input from upstream Time Reducer: Lane = m![N], Time = m![M], Packet = m![P].
    input: ContractTimeTensor<'l, T, f32, m![1], m![1 # 2], m![1 # 256], m![N], m![M], m![P]>,
    // Output: OutTime = m![M, P] = [Time, Packet], OutPacket = m![N] = [Lane].
) -> ContractTensor<'l, T, f32, m![1], m![1 # 2], m![1 # 256], m![M, P], m![N]> {
    input.contract_lane::<m![M, P], m![N]>(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![M], m![P]> = CollectTensor::new(&mut device.main, Tensor::zero());
let b: TrfTensor<bf16, m![1], m![1 # 2], m![1 # 256], m![N], m![P]> = TrfTensor::zero();

let i: ContractTimeTensor<'_, _, f32, m![1], m![1 # 2], m![1 # 256], m![N], m![M], m![P]> = a 
    .contract_outer::<m![M], m![P], m![N], m![P], _>(&b)
    .contract_packet::<m![P]>()
    .contract_time::<m![M]>();

let _o = lane_interleaved(i);
}

Sequential

Lane 차원은 OutTime으로 접혀 들어갑니다. 각 사이클은 한 레인의 Packet에서 8개의 열 위치를 읽으며(플릿당 값 8개), Lane은 연속된 사이클에 걸쳐 순회합니다. 각 사이클의 폭이 8개 요소이므로, Packet은 먼저 8(버스 폭)의 배수로 패딩된 뒤 레인당 PadPacket / 8개 사이클과 사이클당 PadPacket % 8개 요소로 분할됩니다.

PadPacket = Packet # align_up(Packet::SIZE, 8)   (pad Packet up to the next multiple of 8)
OutTime   = [Time, Lane, PadPacket / 8]
OutPacket = [PadPacket % 8]

Packet::SIZE < 32인 경우 [PadPacket / 8]::SIZE = ceil(Packet::SIZE / 8)가 패킷당 사이클 수입니다(예: Packet::SIZE = 4이면 1사이클, Packet::SIZE = 16이면 2사이클).

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![N = 8, M = 4, P = 16];

/// Lane folds into OutTime.
fn lane_sequential<'l, const T: Tu>(
    // Input from upstream Time Reducer: Lane = m![N], Time = m![M], Packet = m![P].
    input: ContractTimeTensor<'l, T, f32, m![1], m![1 # 2], m![1 # 256], m![N], m![M], m![P]>,
    // Output: OutTime = m![M, N, P / 8] = [Time, Lane, Packet / 8], OutPacket = m![P % 8] = [Packet % 8].
) -> ContractTensor<'l, T, f32, m![1], m![1 # 2], m![1 # 256], m![M, N, P / 8], m![P % 8]> {
    input.contract_lane::<m![M, N, P / 8], m![P % 8]>(LaneMode::Sequential)
}

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

let a: CollectTensor<'_, _, bf16, m![1], m![1 # 2], m![1 # 256], m![M], m![P]> = CollectTensor::new(&mut device.main, Tensor::zero());
let b: TrfTensor<bf16, m![1], m![1 # 2], m![1 # 256], m![N], m![P]> = TrfTensor::zero();

let i: ContractTimeTensor<'_, _, f32, m![1], m![1 # 2], m![1 # 256], m![N], m![M], m![P]> = a 
    .contract_outer::<m![M], m![P], m![N], m![P], _>(&b)
    .contract_packet::<m![P]>()
    .contract_time::<m![M]>();

let _o = lane_sequential(i);
}

제약 조건

Lane Folder 자체에는 제약 조건이 없습니다. 여기에서 선택한 LaneMode에 따라 상류 Time Reducer가 적용하는 슬롯 용량 한도가 결정됩니다(Time Reducer 제약 조건 참조).

성능

Interleaved 모드에서 Lane < 8이면 처리량이 Lane::SIZE / 8만큼 감소합니다(비활성 레인이 버스 위치를 비워 두기 때문입니다).

Sequential 모드에서 Packet::SIZE < 8이면(예: Packet Reducer가 bf16 패킷의 절반을 축소한 뒤의 Packet::SIZE = 4) 각 사이클은 8개 요소의 전체 버스 폭이 아니라 정확히 Packet::SIZE개의 요소를 전달합니다. 사이클당 출력은 더 좁아지지만 패딩은 없으며 낭비되는 버스 슬롯도 없습니다.

지연 시간은 무시할 수 있는 수준입니다. Lane Folder는 레인별 출력의 형태만 바꿀 뿐이며, 버퍼를 비우는 데 걸리는 시간 외에 사이클을 추가하지 않습니다.