834 lines
30 KiB
Rust
834 lines
30 KiB
Rust
//! Runtime-generated Compiled Metal expressions from the pinned
|
|
//! metal/compiled.cpp emitter. The shared functor/header source is unchanged.
|
|
use super::array::{Array, Dtype};
|
|
use super::{compiled, ops::Operation};
|
|
use sha2::{Digest, Sha256};
|
|
use std::fmt::Write;
|
|
|
|
#[derive(Debug)]
|
|
struct Input {
|
|
dtype: Dtype,
|
|
scalar: bool,
|
|
constant: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub(super) struct Statement {
|
|
pub(super) dtype: Dtype,
|
|
/// Reference primitive name; Broadcast/AsType emit a static cast.
|
|
pub(super) operator: &'static str,
|
|
pub(super) inputs: Vec<usize>,
|
|
}
|
|
|
|
pub(super) struct Expression {
|
|
inputs: Vec<Input>,
|
|
statements: Vec<Statement>,
|
|
outputs: Vec<usize>,
|
|
prefix: String,
|
|
}
|
|
|
|
pub(super) fn operator(op: &Operation) -> Option<&'static str> {
|
|
use super::{binary::Binary as B, unary::Unary as U};
|
|
Some(match op {
|
|
Operation::Broadcast => "Broadcast",
|
|
Operation::AsType => "AsType",
|
|
Operation::Select => "Select",
|
|
Operation::Unary(op) => match op {
|
|
U::LogicalNot => "LogicalNot",
|
|
U::Abs => "Abs",
|
|
U::Sign => "Sign",
|
|
U::Negative => "Negative",
|
|
U::Square => "Square",
|
|
U::Exp => "Exp",
|
|
U::Log => "Log",
|
|
U::Sigmoid => "Sigmoid",
|
|
U::Tanh => "Tanh",
|
|
U::Sqrt => "Sqrt",
|
|
U::Rsqrt => "Rsqrt",
|
|
U::Sin => "Sin",
|
|
U::Cos => "Cos",
|
|
U::Erf => "Erf",
|
|
},
|
|
Operation::Binary(op) => match op {
|
|
B::Equal => "Equal",
|
|
B::NotEqual => "NotEqual",
|
|
B::GreaterEqual => "GreaterEqual",
|
|
B::Greater => "Greater",
|
|
B::Less => "Less",
|
|
B::LessEqual => "LessEqual",
|
|
B::LogicalAnd => "LogicalAnd",
|
|
B::LogicalOr => "LogicalOr",
|
|
B::BitwiseAnd => "BitwiseAnd",
|
|
B::BitwiseOr => "BitwiseOr",
|
|
B::BitwiseXor => "BitwiseXor",
|
|
B::Remainder => "Remainder",
|
|
B::Add => "Add",
|
|
B::Subtract => "Subtract",
|
|
B::Multiply => "Multiply",
|
|
B::Divide => "Divide",
|
|
B::Maximum => "Maximum",
|
|
B::Minimum => "Minimum",
|
|
B::LogAddExp => "LogAddExp",
|
|
B::Power => "Power",
|
|
},
|
|
_ => return None,
|
|
})
|
|
}
|
|
|
|
fn dtype(d: Dtype) -> &'static str {
|
|
match d {
|
|
Dtype::Bool => "bool",
|
|
Dtype::U8 => "uint8_t",
|
|
Dtype::U16 => "uint16_t",
|
|
Dtype::U32 => "uint32_t",
|
|
Dtype::U64 => "uint64_t",
|
|
Dtype::I8 => "int8_t",
|
|
Dtype::I16 => "int16_t",
|
|
Dtype::I32 => "int32_t",
|
|
Dtype::I64 => "int64_t",
|
|
Dtype::F16 => "float16_t",
|
|
Dtype::BF16 => "bfloat16_t",
|
|
Dtype::F32 => "float",
|
|
Dtype::F64 => "double",
|
|
Dtype::C64 => "complex64_t",
|
|
}
|
|
}
|
|
|
|
// ostream defaultfloat with digits10+1 (7 for F32/F16/BF16, 16 for F64).
|
|
fn float_constant(value: f64, precision: usize) -> String {
|
|
if value.is_nan() {
|
|
return "NAN".into();
|
|
}
|
|
if value.is_infinite() {
|
|
return if value.is_sign_negative() {
|
|
"-INFINITY"
|
|
} else {
|
|
"INFINITY"
|
|
}
|
|
.into();
|
|
}
|
|
let scientific = format!("{:.*e}", precision - 1, value);
|
|
let (mantissa, exponent) = scientific.split_once('e').unwrap();
|
|
let exponent: i32 = exponent.parse().unwrap();
|
|
let trim = |s: String| {
|
|
if s.contains('.') {
|
|
s.trim_end_matches('0').trim_end_matches('.').to_owned()
|
|
} else {
|
|
s
|
|
}
|
|
};
|
|
if exponent < -4 || exponent >= precision as i32 {
|
|
format!("{}e{:+03}", trim(mantissa.to_owned()), exponent)
|
|
} else {
|
|
trim(format!(
|
|
"{:.*}",
|
|
(precision as i32 - 1 - exponent).max(0) as usize,
|
|
value
|
|
))
|
|
}
|
|
}
|
|
|
|
fn constant(a: &Array) -> Result<String, String> {
|
|
a.wait()?;
|
|
if !a.has_data() {
|
|
return Err("compiled constant is unavailable".into());
|
|
}
|
|
let mut bits = [0u8; 8];
|
|
a.buffer().read(a.offset(), &mut bits[..a.itemsize()])?;
|
|
let u16v = u16::from_le_bytes(bits[..2].try_into().unwrap());
|
|
let u32v = u32::from_le_bytes(bits[..4].try_into().unwrap());
|
|
Ok(match a.layout().dtype() {
|
|
Dtype::Bool => (bits[0] != 0).to_string(),
|
|
Dtype::U8 => bits[0].to_string(),
|
|
Dtype::I8 => (bits[0] as i8).to_string(),
|
|
Dtype::U16 => u16v.to_string(),
|
|
Dtype::I16 => (u16v as i16).to_string(),
|
|
Dtype::U32 => u32v.to_string(),
|
|
Dtype::I32 => (u32v as i32).to_string(),
|
|
Dtype::U64 => u64::from_le_bytes(bits).to_string(),
|
|
Dtype::I64 => i64::from_le_bytes(bits).to_string(),
|
|
Dtype::BF16 => float_constant(f32::from_bits(u32::from(u16v) << 16) as f64, 7),
|
|
Dtype::F16 => float_constant(super::super::half_to_f32(u16v) as f64, 7),
|
|
Dtype::F32 => float_constant(f32::from_bits(u32v) as f64, 7),
|
|
Dtype::F64 => float_constant(f64::from_le_bytes(bits), 16),
|
|
Dtype::C64 => format!(
|
|
"complex64_t({}, {})",
|
|
float_constant(f32::from_bits(u32v) as f64, 7),
|
|
float_constant(f32::from_le_bytes(bits[4..].try_into().unwrap()) as f64, 7)
|
|
),
|
|
})
|
|
}
|
|
|
|
struct Names {
|
|
slots: Vec<Option<String>>,
|
|
count: usize,
|
|
}
|
|
impl Names {
|
|
fn get(&mut self, slot: usize) -> String {
|
|
self.slots[slot]
|
|
.get_or_insert_with(|| {
|
|
self.count += 1;
|
|
let mut n = self.count;
|
|
let mut letters = Vec::new();
|
|
while n > 0 {
|
|
letters.push(b'A' + ((n - 1) % 26) as u8);
|
|
n = (n - 1) / 26;
|
|
}
|
|
letters.reverse();
|
|
String::from_utf8(letters).unwrap()
|
|
})
|
|
.clone()
|
|
}
|
|
}
|
|
|
|
impl Expression {
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
pub(super) fn name(&self) -> String {
|
|
std::iter::once("Compiled")
|
|
.chain(self.statements.iter().map(|s| s.operator))
|
|
.collect()
|
|
}
|
|
|
|
pub(super) fn new(
|
|
inputs: &[Array],
|
|
statements: Vec<Statement>,
|
|
outputs: Vec<usize>,
|
|
constants: &[bool],
|
|
) -> Result<Self, String> {
|
|
if inputs.len() != constants.len() || statements.is_empty() || outputs.is_empty() {
|
|
return Err("invalid compiled expression arity".into());
|
|
}
|
|
for (i, s) in statements.iter().enumerate() {
|
|
if s.inputs.is_empty() || s.inputs.iter().any(|&j| j >= inputs.len() + i) {
|
|
return Err("compiled expression is not topological".into());
|
|
}
|
|
}
|
|
if outputs
|
|
.iter()
|
|
.any(|&o| o < inputs.len() || o >= inputs.len() + statements.len())
|
|
{
|
|
return Err("invalid compiled expression output".into());
|
|
}
|
|
let inputs = inputs
|
|
.iter()
|
|
.zip(constants)
|
|
.map(|(a, &embed)| {
|
|
Ok(Input {
|
|
dtype: a.layout().dtype(),
|
|
scalar: a.layout().shape().is_empty(),
|
|
constant: if embed { Some(constant(a)?) } else { None },
|
|
})
|
|
})
|
|
.collect::<Result<Vec<_>, String>>()?;
|
|
// Symbol renaming only; expression spelling and argument ordering
|
|
// follow the original emitter. Hash includes embedded constant text.
|
|
let hash = Sha256::digest(format!("{inputs:?}{statements:?}{outputs:?}"));
|
|
let hash = hash.iter().map(|b| format!("{b:02x}")).collect::<String>();
|
|
let prefix = format!("kernel_qwen_mtplx_compiled_{hash}");
|
|
Ok(Self {
|
|
inputs,
|
|
statements,
|
|
outputs,
|
|
prefix,
|
|
})
|
|
}
|
|
|
|
pub(super) fn evaluate(&self, inputs: &[Array], outputs: &[Array]) -> Result<(), String> {
|
|
super::compile_custom(&format!("{}_contiguous", self.prefix), || self.source())?;
|
|
compiled::evaluate_many(inputs, outputs, &self.prefix, |i| {
|
|
self.inputs[i].constant.is_some()
|
|
})
|
|
}
|
|
|
|
fn value_dtype(&self, slot: usize) -> Dtype {
|
|
if slot < self.inputs.len() {
|
|
self.inputs[slot].dtype
|
|
} else {
|
|
self.statements[slot - self.inputs.len()].dtype
|
|
}
|
|
}
|
|
|
|
pub(super) fn source(&self) -> Result<String, String> {
|
|
let mut source =
|
|
include_str!("../../../../metal/mtplx-qsa-compiled-header.metal").to_owned();
|
|
let work = (8 / self.value_dtype(self.outputs[0]).itemsize()).max(1);
|
|
source.push_str(&self.kernel(true, 0, false, false, 1)?);
|
|
if work > 1 {
|
|
source.push_str(&self.kernel(true, 0, false, false, work)?);
|
|
}
|
|
source.push_str(&self.kernel(true, 0, false, true, work)?);
|
|
for rank in 1..8 {
|
|
source.push_str(&self.kernel(
|
|
false,
|
|
rank,
|
|
false,
|
|
false,
|
|
if rank > 3 { 2 } else { 1 },
|
|
)?);
|
|
source.push_str(&self.kernel(
|
|
false,
|
|
rank,
|
|
false,
|
|
true,
|
|
if rank > 3 { 4 } else { 1 },
|
|
)?);
|
|
}
|
|
source.push_str(&self.kernel(false, 0, true, false, 2)?);
|
|
source.push_str(&self.kernel(false, 0, true, true, 4)?);
|
|
Ok(source)
|
|
}
|
|
|
|
fn kernel(
|
|
&self,
|
|
contiguous: bool,
|
|
rank: usize,
|
|
dynamic: bool,
|
|
large: bool,
|
|
work: usize,
|
|
) -> Result<String, String> {
|
|
let suffix = if contiguous {
|
|
format!("contiguous{}", if work > 1 && !large { "_n" } else { "" })
|
|
} else {
|
|
format!(
|
|
"strided_{}",
|
|
if dynamic {
|
|
"dynamic".to_owned()
|
|
} else {
|
|
rank.to_string()
|
|
}
|
|
)
|
|
};
|
|
let name = format!(
|
|
"{}_{}{}",
|
|
self.prefix,
|
|
suffix,
|
|
if large { "_large" } else { "" }
|
|
);
|
|
let mut os = format!("[[host_name(\"{name}\")]]\n[[kernel]] void {name}(\n");
|
|
let mut names = Names {
|
|
slots: vec![None; self.inputs.len() + self.statements.len()],
|
|
count: 0,
|
|
};
|
|
let mut count = 0;
|
|
let mut indices = false;
|
|
for (i, input) in self.inputs.iter().enumerate() {
|
|
if input.constant.is_some() {
|
|
continue;
|
|
}
|
|
if !input.scalar && !contiguous {
|
|
indices = true;
|
|
}
|
|
writeln!(
|
|
os,
|
|
" device const {}* {} [[buffer({count})]],",
|
|
dtype(input.dtype),
|
|
names.get(i)
|
|
)
|
|
.unwrap();
|
|
count += 1;
|
|
}
|
|
let index_type = if large { "int64_t" } else { "uint" };
|
|
if indices {
|
|
writeln!(
|
|
os,
|
|
" constant const int64_t* in_strides [[buffer({count})]],"
|
|
)
|
|
.unwrap();
|
|
count += 1;
|
|
}
|
|
for &o in &self.outputs {
|
|
writeln!(
|
|
os,
|
|
" device {}* {} [[buffer({count})]],",
|
|
dtype(self.value_dtype(o)),
|
|
names.get(o)
|
|
)
|
|
.unwrap();
|
|
count += 1;
|
|
}
|
|
if contiguous {
|
|
writeln!(
|
|
os,
|
|
" constant const {index_type}& size [[buffer({count})]],"
|
|
)
|
|
.unwrap();
|
|
} else {
|
|
writeln!(
|
|
os,
|
|
" constant const int* output_shape [[buffer({count})]],"
|
|
)
|
|
.unwrap();
|
|
}
|
|
count += 1;
|
|
if dynamic {
|
|
writeln!(os, " constant const int& ndim [[buffer({count})]],").unwrap();
|
|
count += 1;
|
|
}
|
|
if count > 31 {
|
|
return Err("compiled expression exceeds Metal's 31 argument buffers".into());
|
|
}
|
|
os.push_str(
|
|
" uint3 pos [[thread_position_in_grid]],\n uint3 grid [[threads_per_grid]]) {\n",
|
|
);
|
|
writeln!(os, " constexpr int N_ = {work};").unwrap();
|
|
if contiguous && large {
|
|
os.push_str(" int64_t index = N_ * (pos.x + grid.x * int64_t(pos.y));\n");
|
|
} else if contiguous {
|
|
os.push_str(" uint index = N_ * pos.x;\n");
|
|
} else if work > 1 {
|
|
writeln!(
|
|
os,
|
|
" int xshape = output_shape[{}];",
|
|
if dynamic {
|
|
"ndim - 1".to_owned()
|
|
} else {
|
|
(rank - 1).to_string()
|
|
}
|
|
)
|
|
.unwrap();
|
|
writeln!(os, " {index_type} index = N_ * pos.x + xshape * (pos.y + {index_type}(grid.y) * pos.z);").unwrap();
|
|
} else {
|
|
writeln!(
|
|
os,
|
|
" {index_type} index = pos.x + grid.x * (pos.y + {index_type}(grid.y) * pos.z);"
|
|
)
|
|
.unwrap();
|
|
}
|
|
if work > 1 && contiguous {
|
|
os.push_str(" for (int i = 0; i < N_ && index < size; ++i) {\n");
|
|
}
|
|
let mut nc = Vec::new();
|
|
for (i, input) in self.inputs.iter().enumerate() {
|
|
let n = names.get(i);
|
|
if let Some(c) = &input.constant {
|
|
writeln!(
|
|
os,
|
|
" auto tmp_{n} = static_cast<{}>({c});",
|
|
dtype(input.dtype)
|
|
)
|
|
.unwrap();
|
|
} else if input.scalar {
|
|
writeln!(os, " {} tmp_{n} = {n}[0];", dtype(input.dtype)).unwrap();
|
|
} else if contiguous {
|
|
writeln!(os, " {} tmp_{n} = {n}[index];", dtype(input.dtype)).unwrap();
|
|
} else {
|
|
nc.push(i);
|
|
}
|
|
}
|
|
for (i, &slot) in nc.iter().enumerate() {
|
|
let n = names.get(slot);
|
|
write!(os, " {index_type} index_{n} = ").unwrap();
|
|
match rank {
|
|
1 => writeln!(os, "elem_to_loc_1<{index_type}>(pos.x, in_strides[{}]);", i * rank).unwrap(),
|
|
2 => writeln!(os, "elem_to_loc_2<{index_type}>({{pos.x, pos.y}}, in_strides + {});", i * rank).unwrap(),
|
|
3 => writeln!(os, "elem_to_loc_3<{index_type}>(pos, in_strides + {});", i * rank).unwrap(),
|
|
_ if !dynamic => writeln!(os, "N_ * pos.x * {index_type}(in_strides[{}]) + pos.y * {index_type}(in_strides[{}]);", (i+1)*rank-1, (i+1)*rank-2).unwrap(),
|
|
_ => writeln!(os, "N_ * pos.x * {index_type}(in_strides[ndim * {i} + ndim - 1]) + pos.y * {index_type}(in_strides[ndim * {i} + ndim - 2]);").unwrap(),
|
|
}
|
|
}
|
|
if !nc.is_empty() && (rank > 3 || dynamic) {
|
|
os.push_str(" uint zpos = pos.z;\n");
|
|
writeln!(
|
|
os,
|
|
" for (int d = {}; d >= 0; --d) {{",
|
|
if dynamic {
|
|
"ndim - 3".to_owned()
|
|
} else {
|
|
(rank - 3).to_string()
|
|
}
|
|
)
|
|
.unwrap();
|
|
os.push_str(" uint l = zpos % output_shape[d];\n");
|
|
for (i, &slot) in nc.iter().enumerate() {
|
|
let n = names.get(slot);
|
|
if dynamic {
|
|
writeln!(
|
|
os,
|
|
" index_{n} += l * {index_type}(in_strides[{i} * ndim + d]);"
|
|
)
|
|
.unwrap();
|
|
} else {
|
|
writeln!(
|
|
os,
|
|
" index_{n} += l * {index_type}(in_strides[{} + d]);",
|
|
i * rank
|
|
)
|
|
.unwrap();
|
|
}
|
|
}
|
|
os.push_str(" zpos /= output_shape[d];\n }\n");
|
|
}
|
|
if work > 1 && !contiguous {
|
|
os.push_str(" for (int i = 0; i < N_ && (int(N_ * pos.x) + i) < xshape; ++i) {\n");
|
|
}
|
|
for &slot in &nc {
|
|
let n = names.get(slot);
|
|
writeln!(
|
|
os,
|
|
" {} tmp_{n} = {n}[index_{n}];",
|
|
dtype(self.inputs[slot].dtype)
|
|
)
|
|
.unwrap();
|
|
}
|
|
for (i, s) in self.statements.iter().enumerate() {
|
|
let n = names.get(self.inputs.len() + i);
|
|
write!(os, " {} tmp_{n} = ", dtype(s.dtype)).unwrap();
|
|
if matches!(s.operator, "Broadcast" | "AsType") {
|
|
writeln!(
|
|
os,
|
|
"cast_to<{}>(tmp_{});",
|
|
dtype(s.dtype),
|
|
names.get(s.inputs[0])
|
|
)
|
|
.unwrap();
|
|
} else {
|
|
let args = s
|
|
.inputs
|
|
.iter()
|
|
.map(|&s| format!("tmp_{}", names.get(s)))
|
|
.collect::<Vec<_>>()
|
|
.join(", ");
|
|
writeln!(os, "{}()({args});", s.operator).unwrap();
|
|
}
|
|
}
|
|
for &slot in &self.outputs {
|
|
let n = names.get(slot);
|
|
writeln!(os, " {n}[index] = tmp_{n};").unwrap();
|
|
}
|
|
if work > 1 {
|
|
for (i, &slot) in nc.iter().enumerate() {
|
|
let n = names.get(slot);
|
|
if dynamic {
|
|
writeln!(os, " index_{n} += in_strides[{i} * ndim + ndim - 1];").unwrap();
|
|
} else {
|
|
writeln!(os, " index_{n} += in_strides[{}];", i * rank + rank - 1).unwrap();
|
|
}
|
|
}
|
|
os.push_str(" index++;\n }\n");
|
|
}
|
|
os.push_str("}\n");
|
|
Ok(os)
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn mtplx_compiled_constant_spelling() {
|
|
for (value, text) in [
|
|
(0., "0"),
|
|
(-0., "-0"),
|
|
(0.1, "0.1"),
|
|
(1e-4, "0.0001"),
|
|
(1e-5, "1e-05"),
|
|
(1e7, "1e+07"),
|
|
(1234567., "1234567"),
|
|
(12345678., "1.234568e+07"),
|
|
(f64::NAN, "NAN"),
|
|
(f64::INFINITY, "INFINITY"),
|
|
(f64::NEG_INFINITY, "-INFINITY"),
|
|
] {
|
|
assert_eq!(float_constant(value, 7), text);
|
|
}
|
|
assert_eq!(float_constant(1.2345678901234567, 16), "1.234567890123457");
|
|
}
|
|
|
|
#[test]
|
|
fn mtplx_compiled_emitter_matches_all_pinned_silu_bodies() {
|
|
// The pinned shapeless SiLU has two explicit Broadcast primitives;
|
|
// compare the emitter, independently of shape-fixed trace selection.
|
|
let input = Array::unallocated(&[8], Dtype::BF16).unwrap();
|
|
let s = |operator, inputs| Statement {
|
|
dtype: Dtype::BF16,
|
|
operator,
|
|
inputs,
|
|
};
|
|
let expression = Expression::new(
|
|
&[input],
|
|
vec![
|
|
s("Sigmoid", vec![0]),
|
|
s("Broadcast", vec![0, 1]),
|
|
s("Broadcast", vec![1, 0]),
|
|
s("Multiply", vec![2, 3]),
|
|
],
|
|
vec![4],
|
|
&[false],
|
|
)
|
|
.unwrap();
|
|
let source = expression.source().unwrap();
|
|
let reference = include_str!("../../../../metal/mtplx_qwen.metal");
|
|
let mut compared = 0;
|
|
for section in reference
|
|
.split("[[host_name(\"kernel_qwen_mtplx_silu_bf16_")
|
|
.skip(1)
|
|
{
|
|
let suffix = section.split_once('"').unwrap().0;
|
|
let original = format!("[[host_name(\"kernel_qwen_mtplx_silu_bf16_{section}");
|
|
let original = &original[..original.find("\n}\n").unwrap() + 3];
|
|
let symbol = original
|
|
.split_once("void ")
|
|
.unwrap()
|
|
.1
|
|
.split_once('(')
|
|
.unwrap()
|
|
.0;
|
|
let expected = original
|
|
.replace(symbol, &format!("{}_{suffix}", expression.prefix))
|
|
.replace(
|
|
&format!("kernel_qwen_mtplx_silu_bf16_{suffix}"),
|
|
&format!("{}_{suffix}", expression.prefix),
|
|
);
|
|
let marker = format!("[[host_name(\"{}_{suffix}\")]]", expression.prefix);
|
|
let start = source.find(&marker).unwrap();
|
|
let generated = &source[start..];
|
|
let generated = &generated[..generated.find("\n}\n").unwrap() + 3];
|
|
assert_eq!(generated, expected, "{suffix}");
|
|
compared += 1;
|
|
}
|
|
assert_eq!(compared, 19);
|
|
}
|
|
|
|
#[test]
|
|
fn mtplx_compile_fusion_preserves_parent_and_output_boundaries() {
|
|
use super::{
|
|
binary::{Binary, binary},
|
|
compiled::Replay,
|
|
ops,
|
|
stream::{Device, Stream},
|
|
unary::{Unary, unary},
|
|
};
|
|
let gpu = Stream::new(0, Device::Gpu);
|
|
let other = Stream::new(1, Device::Gpu);
|
|
let x = Array::unallocated(&[4], Dtype::F32).unwrap();
|
|
let scalar = Array::unallocated(&[], Dtype::F32).unwrap();
|
|
let inputs = [x, scalar];
|
|
let results = Replay::trace_compiled(&inputs, |a| {
|
|
let square = unary(&a[0], Unary::Square, gpu)?;
|
|
let sum = binary(&square, &a[1], Binary::Add, gpu)?;
|
|
Ok(vec![square, sum.clone(), sum])
|
|
})
|
|
.unwrap()
|
|
.run(&inputs)
|
|
.unwrap();
|
|
assert!(matches!(
|
|
&*results[0].operation().unwrap(),
|
|
Operation::Compiled(_)
|
|
));
|
|
assert_eq!(results[0].outputs().len(), 2);
|
|
assert_eq!(results[1].id(), results[2].id());
|
|
assert!(std::rc::Rc::ptr_eq(
|
|
&results[0].operation().unwrap(),
|
|
&results[1].operation().unwrap()
|
|
));
|
|
// Mixed-parent Broadcast is split; the outside consumer retains its view.
|
|
let results = Replay::trace_compiled(&inputs, |a| {
|
|
let broadcast = ops::broadcast_to(&a[1], &[4], gpu)?;
|
|
let sum = binary(&a[0], &broadcast, Binary::Add, gpu)?;
|
|
Ok(vec![ops::reshape(&broadcast, &[2, 2], gpu)?, sum])
|
|
})
|
|
.unwrap()
|
|
.run(&inputs)
|
|
.unwrap();
|
|
assert!(matches!(
|
|
&*results[0].inputs()[0].operation().unwrap(),
|
|
Operation::Broadcast
|
|
));
|
|
assert!(results[1].inputs().iter().any(|a| a.id() == inputs[1].id()));
|
|
// A non-Broadcast intermediate with outside parents cannot be split.
|
|
let results = Replay::trace_compiled(&inputs, |a| {
|
|
let square = unary(&a[0], Unary::Square, gpu)?;
|
|
Ok(vec![
|
|
ops::reshape(&square, &[2, 2], gpu)?,
|
|
binary(&square, &a[1], Binary::Add, gpu)?,
|
|
])
|
|
})
|
|
.unwrap()
|
|
.run(&inputs)
|
|
.unwrap();
|
|
assert_eq!(results[0].inputs()[0].id(), results[1].inputs()[0].id());
|
|
// Stream mismatch is a fusion boundary even though CSE ignores streams.
|
|
let results = Replay::trace_compiled(&inputs, |a| {
|
|
let square = unary(&a[0], Unary::Square, gpu)?;
|
|
let negative = unary(&square, Unary::Negative, other)?;
|
|
Ok(vec![binary(&negative, &a[1], Binary::Add, other)?])
|
|
})
|
|
.unwrap()
|
|
.run(&inputs)
|
|
.unwrap();
|
|
assert_eq!(results[0].inputs()[0].stream(), Some(gpu));
|
|
// A 25-node chain splits at the original depth=11 boundary: 11+11+3.
|
|
let results = Replay::trace_compiled(&inputs, |a| {
|
|
let mut x = a[0].clone();
|
|
for _ in 0..25 {
|
|
x = unary(&x, Unary::Negative, gpu)?;
|
|
}
|
|
Ok(vec![x])
|
|
})
|
|
.unwrap()
|
|
.run(&inputs)
|
|
.unwrap();
|
|
let mut x = results[0].clone();
|
|
let mut lengths = Vec::new();
|
|
while let Some(op) = x.operation() {
|
|
let Operation::Compiled(expression) = &*op else {
|
|
panic!("unfused chain");
|
|
};
|
|
lengths.push(expression.statements.len());
|
|
let next = x.inputs()[0].clone();
|
|
x = next;
|
|
}
|
|
assert_eq!(lengths, [11, 11, 3]);
|
|
// Broadcast-only outputs remain views rather than materialized kernels.
|
|
let result = Replay::trace_compiled(&inputs, |a| {
|
|
Ok(vec![ops::broadcast_to(&a[1], &[4096], gpu)?])
|
|
})
|
|
.unwrap()
|
|
.run(&inputs)
|
|
.unwrap();
|
|
assert!(matches!(
|
|
&*result[0].operation().unwrap(),
|
|
Operation::Broadcast
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires Metal and external supervisor; generated fusion vs pinned SiLU plus multi-output/scalar dispatch"]
|
|
fn mtplx_compiled_fusion_gpu_matches_reference() {
|
|
use super::{
|
|
super::{
|
|
gpu::Context,
|
|
qwen_mtplx_tests::{capture_dispatches, pattern},
|
|
},
|
|
allocator,
|
|
compiled::Replay,
|
|
configure_sources, gdn, mlp, ops,
|
|
stream::{Device, Streams},
|
|
};
|
|
use super::{
|
|
binary::{Binary, binary},
|
|
unary::{Unary, unary},
|
|
};
|
|
configure_sources().unwrap();
|
|
let _context = Context::open_qwen(0).unwrap();
|
|
let streams = Streams::new(Some(allocator::Allocator::new().unwrap()));
|
|
let gpu = streams.default_stream(Device::Gpu).unwrap();
|
|
for rank in [1, 4, 8] {
|
|
for reverse in [false, true] {
|
|
eprintln!("compiled SiLU rank={rank} reverse={reverse}");
|
|
let shape = vec![2; rank];
|
|
let input = Array::new(&shape, Dtype::BF16, pattern(1 << rank, 37)).unwrap();
|
|
let axes = (0..rank as i32)
|
|
.step_by(2)
|
|
.chain((1..rank as i32).step_by(2))
|
|
.collect::<Vec<_>>();
|
|
let input = ops::transpose(&input, &axes, gpu).unwrap();
|
|
let input = if reverse {
|
|
ops::slice(
|
|
&input,
|
|
&vec![1; rank],
|
|
&vec![-3; rank],
|
|
&vec![-1; rank],
|
|
gpu,
|
|
)
|
|
.unwrap()
|
|
} else {
|
|
input
|
|
};
|
|
ops::evaluate(&streams, std::slice::from_ref(&input), gpu, false).unwrap();
|
|
let expected = mlp::silu(&input, gpu).unwrap();
|
|
let graph = Replay::trace_compiled(std::slice::from_ref(&input), |a| {
|
|
Ok(vec![mlp::silu(&a[0], gpu)?])
|
|
})
|
|
.unwrap();
|
|
let actual = graph.run(std::slice::from_ref(&input)).unwrap().remove(0);
|
|
let (_, calls) = capture_dispatches(|| {
|
|
ops::evaluate(&streams, std::slice::from_ref(&actual), gpu, false).unwrap()
|
|
});
|
|
assert_eq!(calls.len(), 1);
|
|
assert!(calls[0].0.starts_with("kernel_qwen_mtplx_compiled_"));
|
|
let actual = ops::contiguous(&actual, false, gpu).unwrap();
|
|
let expected = ops::contiguous(&expected, false, gpu).unwrap();
|
|
ops::evaluate(&streams, &[actual.clone(), expected.clone()], gpu, false).unwrap();
|
|
let mut a = vec![0; actual.layout().nbytes()];
|
|
let mut b = vec![0; expected.layout().nbytes()];
|
|
actual.buffer().read(actual.offset(), &mut a).unwrap();
|
|
expected.buffer().read(expected.offset(), &mut b).unwrap();
|
|
assert_eq!(a, b);
|
|
}
|
|
}
|
|
for rows in [1, 4, 17] {
|
|
eprintln!("compiled nested compute_g rows={rows}");
|
|
let args = [
|
|
Array::new(&[48], Dtype::BF16, pattern(48, 62)).unwrap(),
|
|
Array::new(&[1, rows, 48], Dtype::BF16, pattern(rows as u32 * 48, 60)).unwrap(),
|
|
Array::new(&[48], Dtype::BF16, pattern(48, 63)).unwrap(),
|
|
];
|
|
let expected = gdn::compute_g(&args[0], &args[1], &args[2], gpu).unwrap();
|
|
let graph = Replay::trace_compiled(&args, |a| {
|
|
Ok(vec![gdn::compute_g(&a[0], &a[1], &a[2], gpu)?])
|
|
})
|
|
.unwrap();
|
|
let actual = graph.run(&args).unwrap().remove(0);
|
|
let (_, calls) = capture_dispatches(|| {
|
|
ops::evaluate(&streams, std::slice::from_ref(&actual), gpu, false).unwrap()
|
|
});
|
|
assert_eq!(calls.len(), 1, "{calls:?}");
|
|
assert!(calls[0].0.starts_with("kernel_qwen_mtplx_compiled_"));
|
|
ops::evaluate(&streams, std::slice::from_ref(&expected), gpu, false).unwrap();
|
|
let mut a = vec![0; actual.layout().nbytes()];
|
|
let mut b = vec![0; expected.layout().nbytes()];
|
|
actual.buffer().read(actual.offset(), &mut a).unwrap();
|
|
expected.buffer().read(expected.offset(), &mut b).unwrap();
|
|
assert_eq!(a, b, "nested compute_g rows={rows}");
|
|
}
|
|
for n in [8, 65536] {
|
|
for embedded in [false, true] {
|
|
eprintln!("compiled multiple outputs n={n} embedded={embedded}");
|
|
let data = (0..n)
|
|
.flat_map(|i| ((i % 8) as f32).to_le_bytes())
|
|
.collect::<Vec<_>>();
|
|
let input = Array::new(&[n], Dtype::F32, super::scalar_buffer(&data).unwrap()).unwrap();
|
|
let scalar = Array::new(
|
|
&[],
|
|
Dtype::F32,
|
|
super::scalar_buffer(&2f32.to_le_bytes()).unwrap(),
|
|
)
|
|
.unwrap();
|
|
let args = if embedded {
|
|
vec![input]
|
|
} else {
|
|
vec![input, scalar.clone()]
|
|
};
|
|
let graph = Replay::trace_compiled(&args, |a| {
|
|
let square = unary(&a[0], Unary::Square, gpu)?;
|
|
let sum = binary(
|
|
&square,
|
|
if embedded { &scalar } else { &a[1] },
|
|
Binary::Add,
|
|
gpu,
|
|
)?;
|
|
Ok(vec![square, sum])
|
|
})
|
|
.unwrap();
|
|
let roots = graph.run(&args).unwrap();
|
|
// The pinned evaluator selects its pipeline before allocating
|
|
// outputs (metal/compiled.cpp:384,419). Fresh ArrayDesc data_size
|
|
// is zero, even when the logical shape exceeds 65536 elements.
|
|
assert!(roots.iter().all(|a| a.data_size() == 0));
|
|
let (_, calls) =
|
|
capture_dispatches(|| ops::evaluate(&streams, &roots, gpu, false).unwrap());
|
|
assert_eq!(calls.len(), 1, "{calls:?}");
|
|
assert!(calls[0].0.ends_with("_contiguous"), "{calls:?}");
|
|
assert_eq!(calls[0].1, [n as u32, 1, 1]);
|
|
for (j, root) in roots.iter().enumerate() {
|
|
let mut bytes = vec![0; root.layout().nbytes()];
|
|
root.buffer().read(root.offset(), &mut bytes).unwrap();
|
|
for (i, b) in bytes.chunks_exact(4).enumerate() {
|
|
assert_eq!(
|
|
f32::from_le_bytes(b.try_into().unwrap()),
|
|
((i % 8) * (i % 8)) as f32 + (j * 2) as f32
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|