사례 연구: Tensor Unit I/O

이 엔드투엔드 Tensor Unit I/O 패턴은 Fetch와 Commit을 사용해 모든 논리 인덱스를 보존하면서 텐서를 DM을 거쳐 단계적으로 옮깁니다. 이는 여러 동작을 조합한 이동 패턴이며, 개별 Fetch Engine 및 Commit Engine 레퍼런스 페이지와는 구별됩니다.

대표적인 공개 사례는 fetch_commit_simple입니다. 이 사례는 매핑이 [A, B]인 i8 HBM 텐서를 받아 Tensor Unit을 통해 값을 i32로 확장하고, 매핑이 [B, A]인 HBM 텐서를 반환합니다. 마지막 to_hbm DMA 호출이 출력 HBM 매핑을 선택하므로, 이 테스트는 HBM-DM 전송, Fetch, Collect, Commit, DM-HBM 전송으로 이어지는 전체 파이프라인을 검증합니다. 이 사례는 의도적으로 스트리밍 순열을 주장하지 않습니다. 공개 예제에서 [A, B]를 [B, A]로 바꾸는 재배치는 마지막 DMA 경계에서 일어납니다. 디바이스 소스는 아래에 포함되어 있습니다.

#![expect(clippy::type_complexity)]

use furiosa_opt_std::prelude::*;

axes![A = 4096, B = 8];

type Chip = m![1];
type Cluster = m![1 # 2];

#[device(chip = 1)]
pub fn fetch_commit_simple(
    device: &mut Device,
    input: &HbmTensor<i8, m![1], m![A, B]>,
) -> HbmTensor<i32, m![1], m![B, A]> {
    // Element's innermost axis must match the source's innermost (B, stride 1) so that the
    // DMA tail contains the full B axis (8 × i8 = 8 bytes), satisfying min_align = 8.
    let input_dm = input.to_dm::<Cluster, m![A / 16], m![A / 8 % 2, A % 8, B]>(&mut device.tdma);

    let fetch_and_commit_tensor: DmTensor<i32, Chip, Cluster, m![A / 16], m![A / 8 % 2, A % 8, B]> = device
        .main
        .begin(input_dm.view())
        .fetch::<m![A / 8 % 2], m![A % 8, B]>()
        .fetch_cast::<i32>()
        .collect::<m![A / 8 % 2, A % 8], m![B]>()
        .commit_trim::<m![B]>()
        .commit();

    fetch_and_commit_tensor.to_hbm(&mut device.tdma)
}

호스트 오라클과 테스트는 examples 테스트 타깃에 포함되어 있습니다.

use furiosa_opt_examples::fetch_commit::fetch_commit_simple;
use furiosa_opt_std::prelude::*;

#[tokio::test]
async fn test_fetch_commit_simple_host() {
    use furiosa_opt_examples::fetch_commit::{A, B};

    let mut device = Device::new(fetch_commit_simple.topology()).unwrap();

    // Create input tensor with shape (A=4096)(B=8).
    let input = HostTensor::<i8, m![A, B]>::from_vec((0..32768).map(|x| x as i8).collect::<Vec<_>>())
        .to_hbm::<m![1], m![A, B]>(&mut device.pdma)
        .await
        .unwrap();

    // Call the device function.
    let output = launch(fetch_commit_simple, (&mut device, &input)).await.unwrap();

    let mut expected = vec![0i32; 4096 * 8];
    let mut idx = 0;
    for b in 0..8 {
        for a in 0..4096 {
            expected[idx] = ((a * 8 + b) as i8) as i32;
            idx += 1;
        }
    }

    assert_eq!(
        output.to_host::<m![B, A]>(&mut device.pdma).await.unwrap().into_inner(),
        Tensor::<_, m![B, A], CurrentBackend>::from_vec(expected)
    );
}

대표 사례를 실행하려면 cargo furiosa-opt test --test fetch_commit_tests를 실행하십시오. 호스트 오라클은 모든 output[b, a] 값을 확장된 input[a, b] 값과 비교하여 검사합니다. 스트리밍 경로는 입력 DM, Fetch, Cast, Collect, Commit, 그리고 원래의 [A, B] DM 레이아웃을 유지하는 출력 DM으로 이어집니다. 여기서 사용한 API는 Fetch Engine, Collect Engine, Commit Engine 문서에 설명되어 있습니다.