ruda-kernel 0.1.53

Ruda Rust kernel DSL and device tensor operations.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

use crate::dsl::prelude::*;
use alloc::vec;
use half::f16;
use ruda_core::ir::{CoopMma, FloatKind, Operation, PointerOp, Type, UIntKind};

#[ruda]
fn vectors(a: &Vector<f16, Const<8>>, b: &Vector<f16, Const<4>>, c: &mut Vector<f32, Const<4>>) {
    let definition = crate::dsl::cmma::MmaDefinition::<f16, f16, f32>::new(16usize, 8usize, 16usize);
    definition.execute_vectors_inplace(a, b, c);
}

#[ruda]
fn arrays(a: &Array<Vector<f16, Const<2>>>, b: &Array<Vector<f16, Const<2>>>, c: &mut Vector<f32, Const<4>>) {
    let definition = crate::dsl::cmma::MmaDefinition::<f16, f16, f32>::new(16usize, 8usize, 16usize);
    definition.execute_vector_inplace(a, b, c);
}

#[test]
fn manual_mma_vector_sequence_index_mut_keeps_the_accumulator_addressable() {
    let mut scope = Scope::root(false);
    let c = scope.create_local(Type::from(FloatKind::F32).with_vector_size(4));
    let mut sequence = Sequence::<Vector<f32, Const<4>>>::__expand_new(&mut scope);
    sequence.__expand_push_method(&mut scope, c.clone().into());
    let reference = sequence.__expand_index_mut_method(&mut scope, 0);
    assert!(scope.reference_storage(*c).is_some());
    let pointer = scope.reference_pointer(*reference.expand).expect("mutable sequence element reference");
    let a = scope.create_local(Type::from(FloatKind::F16).with_vector_size(8));
    let b = scope.create_local(Type::from(FloatKind::F16).with_vector_size(4));
    scope.register_type::<usize>(UIntKind::U32.into());
    vectors::expand(&mut scope, a.into(), b.into(), reference);
    assert_mma_writeback(&scope, pointer);
}

#[test]
fn manual_mma_vector_intrinsics_read_values_and_write_through_accumulator_reference() {
    for vector_operands in [false, true] {
        let mut scope = Scope::root(false);
        scope.register_type::<usize>(UIntKind::U32.into());
        let half = Type::from(FloatKind::F16);
        let c = scope.create_local(Type::from(FloatKind::F32).with_vector_size(4));
        let c_reference = scope.reference_borrow(c, true).unwrap();
        let c_pointer = scope.reference_pointer(*c_reference).unwrap();
        if vector_operands {
            let a = scope.create_local(half.with_vector_size(8));
            let b = scope.create_local(half.with_vector_size(4));
            let a = scope.reference_borrow(a, false).unwrap();
            let b = scope.reference_borrow(b, false).unwrap();
            vectors::expand(&mut scope, a.into(), b.into(), c_reference.into());
        } else {
            let a = scope.create_local_array(half.with_vector_size(2), 4);
            let b = scope.create_local_array(half.with_vector_size(2), 2);
            arrays::expand(&mut scope, a.into(), b.into(), c_reference.into());
        }
        assert_mma_writeback(&scope, c_pointer);
    }
}

fn assert_mma_writeback(scope: &Scope, pointer: ruda_core::ir::Variable) {
        let (position, result) = scope.instructions.iter().enumerate().find_map(|(i, instruction)| {
            if let Operation::CoopMma(CoopMma::ExecuteManual { registers_a, registers_b, registers_c, .. }) = instruction.operation {
                for value in [registers_a, registers_b, registers_c] {
                    assert!(scope.reference_target(value).is_none(), "MMA must read a value, not a reference proxy");
                }
                Some((i, instruction.out.unwrap()))
            } else { None }
        }).expect("manual MMA instruction");
        let mut pointers = vec![pointer];
        for instruction in &scope.instructions[..position] {
            if let Operation::Copy(input) = instruction.operation {
                if pointers.contains(&input) { pointers.push(instruction.out.unwrap()); }
            }
        }
        let mut values = vec![result];
        let mut writeback = false;
        for instruction in &scope.instructions[position + 1..] {
            match instruction.operation {
                Operation::Copy(input) if values.contains(&input) => values.push(instruction.out.unwrap()),
                Operation::Pointer(PointerOp::Store { pointer, value, .. })
                    if pointers.contains(&pointer) && values.contains(&value) => writeback = true,
                _ => {}
            }
        }
        assert!(writeback, "MMA result must update the caller's accumulator, not a discarded proxy");
}