Cast Engine

Cast Engine은 Commit Engine이 f32/i32 파이프라인 결과를 DM에 쓰기 전에 이를 더 낮은 정밀도 타입(예: bf16)으로 좁혀 저장 비용을 줄입니다.

인터페이스

CollectTensor, ContractTensor, VectorFinalTensor는 모두 동일한 의미의 .cast()를 제공합니다.

//
// The Cast Engine accepts only `VeScalar` inputs (hardware constraint), so the
// bound lives on the impl rather than on a wider trait.
impl<'l, const T: Tu, P: CanApplyCast, D: VeScalar, Chip: M, Cluster: M, Slice: M, Time: M, Packet: M, B: Backend>
    TuTensor<'l, T, P, D, Chip, Cluster, Slice, Time, Packet, B>
{
    /// Casts each element to type `OutD` and pads the output packet back to one
    /// 32-byte flit.
    #[primitive(TuTensor::cast)]
    pub fn cast<OutD: Scalar, OutPacket: M>(self) -> CastTensor<'l, T, OutD, Chip, Cluster, Slice, Time, OutPacket, B>
    where
        D: CastEngineCast<OutD>,
    {
        verify_cast::<D, OutD, Packet, OutPacket>();
        CastTensor::new(self.device, self.inner.map(|v| v.cast()).transpose(false))
    }
}

.cast::<OutD, OutPacket>()는 각 요소를 OutD 타입으로 변환하고 출력을 다시 하나의 32바이트 플릿으로 패딩합니다. 커널 작성자는 OutD(대상 타입)와 OutPacket(출력 요소 레이아웃)을 선택합니다. 나머지는 컴파일러가 도출합니다.

Cast Engine은 수학적 텐서 이동은 아니지만, 텐서의 형태를 보존하고 요소 타입만 변경합니다. Packet 레이아웃을 제외한 모든 차원은 변경 없이 그대로 통과하며, Packet 레이아웃은 출력이 여전히 하나의 32바이트 플릿에 들어가도록 다시 패딩됩니다.

아래 예제는 8개 요소로 이루어진 i32 패킷(8 × 4 = 32바이트)을 i8로 캐스팅합니다. 캐스팅 후 8개 요소는 8바이트를 차지하므로, A # 32가 출력을 다시 32바이트로 패딩합니다:

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

fn cast_i32_to_i8<'l, const T: Tu>(
    input: CollectTensor<'l, T, i32, m![1], m![1 # 2], m![1 # 256], m![B], m![A]>,
) -> CastTensor<'l, T, i8, m![1], m![1 # 2], m![1 # 256], m![B], m![A # 32]> {
    input.cast()
}

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

let c: CollectTensor<'_, _, i32, m![1], m![1 # 2], m![1 # 256], m![B], m![A]> = CollectTensor::new(&mut device.main, Tensor::zero());
let _o = cast_i32_to_i8(c);
}

입력 데이터가 32바이트를 모두 채우지 않을 수도 있습니다. 아래 예제는 4개의 데이터 요소가 8로 패딩된(A # 8, 32바이트) i32 입력을, 동일한 4개 요소가 32로 패딩된(A # 32, 역시 32바이트) i8 출력으로 캐스팅합니다:

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

fn cast_padded<'l, const T: Tu>(
    input: CollectTensor<'l, T, i32, m![1], m![1 # 2], m![1 # 256], m![1], m![A # 8]>,
) -> CastTensor<'l, T, i8, m![1], m![1 # 2], m![1 # 256], m![1], m![A # 32]> {
    input.cast()
}

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

let c: CollectTensor<'_, _, i32, m![1], m![1 # 2], m![1 # 256], m![1], m![A # 8]> = CollectTensor::new(&mut device.main, Tensor::zero());
let _o = cast_padded(c);
}

지원되는 캐스팅

각 입력은 32바이트 플릿입니다. 지원되는 원본 타입은 f32와 i32이며, 각각 지원하는 대상 타입이 정해져 있습니다:

성능

Cast Engine은 결코 파이프라인 병목이 되지 않습니다. 플릿 중 유효한 데이터를 담은 부분이 얼마이든 상관없이 사이클당 하나의 플릿을 처리하기 때문입니다. 다운스트림 Commit Engine이 충분히 활용되지 않은 플릿을 모아 밀집된 DM 쓰기로 만들므로 DM 대역폭이 낭비되지 않습니다.