use fuel_graph::registry::{FusedOpId, FusedOps};
use fuel_ir::DType;
use fuel_ir::dispatch::OpKind;
use fuel_ir::probe::BackendId;
use smallvec::SmallVec;
use crate::fkc::caps_map::{self, ResolvedLayout};
use crate::fkc::cost_expr::{self, CompiledCostExpr};
use crate::fkc::error::FkcError;
use crate::fkc::precision;
use crate::fkc::revhash;
use crate::fkc::schema::{CostBlock, FkcKernel};
use crate::fused::{KernelRevisionHash, PrecisionGuarantee};
use crate::kernel::{CostFn, KernelCaps, KernelDTypes, KernelRef};
pub trait LinkRegistry {
fn resolve_primitive(&self, symbol: &str) -> Option<KernelRef>;
fn resolve_fused(&self, symbol: &str) -> Option<KernelRef>;
fn resolve_cost_fn(&self, _name: &str) -> Option<CostFn> {
None
}
}
#[derive(Debug, Clone)]
pub struct ResolvedPrimitive {
pub op: OpKind,
pub dtypes: KernelDTypes,
pub backend: BackendId,
pub caps: KernelCaps,
pub layouts: Vec<ResolvedLayout>,
pub precision: PrecisionGuarantee,
pub cost: CompiledCostExpr,
pub cost_fn: Option<CostFn>,
pub kernel: KernelRef,
pub kernel_source: String,
pub revision: KernelRevisionHash,
pub variant: Option<String>,
}
#[derive(Debug, Clone)]
pub struct ResolvedFused {
pub id: FusedOpId,
pub dtypes: KernelDTypes,
pub backend: BackendId,
pub caps: KernelCaps,
pub layouts: Vec<ResolvedLayout>,
pub precision: PrecisionGuarantee,
pub cost: CompiledCostExpr,
pub kernel: KernelRef,
pub kernel_source: String,
pub revision: KernelRevisionHash,
pub variant: Option<String>,
pub bundle_slot_names: Vec<String>,
}
#[derive(Debug, Clone)]
pub enum Resolved {
Primitive(ResolvedPrimitive),
Fused(ResolvedFused),
}
pub(crate) fn lower_op_kind(s: &str, section: &str) -> Result<OpKind, FkcError> {
let mapped = match s {
"MatMul" => Some(OpKind::MatMul),
"AddElementwise" => Some(OpKind::AddElementwise),
"SubElementwise" => Some(OpKind::SubElementwise),
"MulElementwise" => Some(OpKind::MulElementwise),
"DivElementwise" => Some(OpKind::DivElementwise),
"ReluElementwise" => Some(OpKind::ReluElementwise),
"NegElementwise" => Some(OpKind::NegElementwise),
"SqrElementwise" => Some(OpKind::SqrElementwise),
"SqrtElementwise" => Some(OpKind::SqrtElementwise),
"RecipElementwise" => Some(OpKind::RecipElementwise),
"AbsElementwise" => Some(OpKind::AbsElementwise),
"TanhElementwise" => Some(OpKind::TanhElementwise),
"ExpElementwise" => Some(OpKind::ExpElementwise),
"LogElementwise" => Some(OpKind::LogElementwise),
"SinElementwise" => Some(OpKind::SinElementwise),
"CosElementwise" => Some(OpKind::CosElementwise),
"SigmoidElementwise" => Some(OpKind::SigmoidElementwise),
"SiluElementwise" => Some(OpKind::SiluElementwise),
"GeluElementwise" => Some(OpKind::GeluTanhElementwise),
"StepElementwise" => Some(OpKind::StepElementwise),
"SumReduce" => Some(OpKind::SumReduce),
"MaxReduce" => Some(OpKind::MaxReduce),
"MinReduce" => Some(OpKind::MinReduce),
"MeanReduce" => Some(OpKind::MeanReduce),
"Cast" => Some(OpKind::Cast),
"Conv2D" => Some(OpKind::Conv2D),
"ConvTranspose2D" => Some(OpKind::ConvTranspose2D),
"ReduceSumTo" => Some(OpKind::ReduceSumTo),
"ReduceMaxTo" => Some(OpKind::ReduceMaxTo),
"FusedLinear" => Some(OpKind::FusedLinear),
"FlashAttn" => Some(OpKind::FlashAttn),
"FlashAttnBackwardQ" => Some(OpKind::FlashAttnBackwardQ),
"FlashAttnBackwardK" => Some(OpKind::FlashAttnBackwardK),
"FlashAttnBackwardV" => Some(OpKind::FlashAttnBackwardV),
"PagedAttn" => Some(OpKind::PagedAttn),
"Affine" => Some(OpKind::Affine),
"ClampElementwise" => Some(OpKind::ClampElementwise),
"PowIElementwise" => Some(OpKind::PowIElementwise),
"PowIElementwiseBackward" => Some(OpKind::PowIElementwiseBackward),
"MaximumElementwise" => Some(OpKind::MaximumElementwise),
"MinimumElementwise" => Some(OpKind::MinimumElementwise),
"EqualElementwise" => Some(OpKind::EqualElementwise),
"NotEqualElementwise" => Some(OpKind::NotEqualElementwise),
"LessElementwise" => Some(OpKind::LessElementwise),
"LessEqualElementwise" => Some(OpKind::LessEqualElementwise),
"GreaterElementwise" => Some(OpKind::GreaterElementwise),
"GreaterEqualElementwise" => Some(OpKind::GreaterEqualElementwise),
"Where" => Some(OpKind::Where),
"FloorElementwise" => Some(OpKind::FloorElementwise),
"CeilElementwise" => Some(OpKind::CeilElementwise),
"RoundElementwise" => Some(OpKind::RoundElementwise),
"SignElementwise" => Some(OpKind::SignElementwise),
"ErfElementwise" => Some(OpKind::ErfElementwise),
"GeluErfElementwise" => Some(OpKind::GeluElementwise),
"PowElementwise" => Some(OpKind::PowElementwise),
"RsqrtElementwise" => Some(OpKind::RsqrtElementwise),
"RemElementwise" => Some(OpKind::RemElementwise),
"Flip" => Some(OpKind::Flip),
"Roll" => Some(OpKind::Roll),
"CumSum" => Some(OpKind::CumSum),
"Pad" => Some(OpKind::Pad),
"Triu" => Some(OpKind::Triu),
"Tril" => Some(OpKind::Tril),
"LogSoftmaxLastDim" => Some(OpKind::LogSoftmaxLastDim),
"LogSoftmaxLastDimBackward" => Some(OpKind::LogSoftmaxLastDimBackward),
"MaskedFill" => Some(OpKind::MaskedFill),
"PadBackward" => Some(OpKind::PadBackward),
"Concat" => Some(OpKind::Concat),
"SoftmaxLastDim" => Some(OpKind::SoftmaxLastDim),
"SoftmaxLastDimBackward" => Some(OpKind::SoftmaxLastDimBackward),
"RmsNormLastDim" => Some(OpKind::RmsNormLastDim),
"RmsNormLastDimBackward" => Some(OpKind::RmsNormLastDimBackward),
"LayerNormLastDim" => Some(OpKind::LayerNormLastDim),
"LayerNormLastDimBackward" => Some(OpKind::LayerNormLastDimBackward),
"ReduceMaxToBackward" => Some(OpKind::ReduceMaxToBackward),
"IndexSelect" => Some(OpKind::IndexSelect),
"Gather" => Some(OpKind::Gather),
"Rope" => Some(OpKind::Rope),
"IndexAdd" => Some(OpKind::IndexAdd),
"ScatterAdd" => Some(OpKind::ScatterAdd),
"ArgMaxDim" => Some(OpKind::ArgMaxDim),
"ArgMinDim" => Some(OpKind::ArgMinDim),
"QMatMul" => Some(OpKind::QMatMul),
"WriteSlice" => Some(OpKind::WriteSlice),
"WriteSliceRotating" => Some(OpKind::WriteSliceRotating),
"WriteSliceDoff" => Some(OpKind::WriteSliceDoff),
"Copy" => Some(OpKind::Copy),
"ReluInplace" => Some(OpKind::ReluInplace),
"SiluInplace" => Some(OpKind::SiluInplace),
"GeluInplace" => Some(OpKind::GeluTanhInplace),
"TanhInplace" => Some(OpKind::TanhInplace),
"SigmoidInplace" => Some(OpKind::SigmoidInplace),
"NegInplace" => Some(OpKind::NegInplace),
"AbsInplace" => Some(OpKind::AbsInplace),
"SqrInplace" => Some(OpKind::SqrInplace),
"SqrtInplace" => Some(OpKind::SqrtInplace),
"RsqrtInplace" => Some(OpKind::RsqrtInplace),
"RecipInplace" => Some(OpKind::RecipInplace),
"ExpInplace" => Some(OpKind::ExpInplace),
"LogInplace" => Some(OpKind::LogInplace),
"SinInplace" => Some(OpKind::SinInplace),
"CosInplace" => Some(OpKind::CosInplace),
"SignInplace" => Some(OpKind::SignInplace),
"FloorInplace" => Some(OpKind::FloorInplace),
"CeilInplace" => Some(OpKind::CeilInplace),
"RoundInplace" => Some(OpKind::RoundInplace),
"ErfInplace" => Some(OpKind::ErfInplace),
"GeluErfInplace" => Some(OpKind::GeluInplace),
"ClampInplace" => Some(OpKind::ClampInplace),
"PowIInplace" => Some(OpKind::PowIInplace),
"InplaceAffine" => Some(OpKind::InplaceAffine),
"FusedSoftmaxCrossEntropy" => Some(OpKind::FusedSoftmaxCrossEntropy),
"CausalConv1d" => Some(OpKind::CausalConv1d),
"SelectiveScan" => Some(OpKind::SelectiveScan),
"SsdChunkScan" => Some(OpKind::SsdChunkScan),
"Nf4Matmul" => Some(OpKind::Nf4Matmul),
_ => None,
};
match mapped {
Some(op) => Ok(op),
None => Err(FkcError::UnknownOpKind {
section: section.to_string(),
op_kind: s.to_string(),
}),
}
}
pub(crate) fn lower_fused_op(s: &str, section: &str) -> Result<FusedOpId, FkcError> {
fused_op_id_for_const_name(s).ok_or_else(|| FkcError::UnknownFusedOp {
section: section.to_string(),
fused_op: s.to_string(),
})
}
pub(crate) fn fused_op_id_for_const_name(s: &str) -> Option<FusedOpId> {
let id = match s {
"SOFTMAX_LAST_DIM" => FusedOps::SOFTMAX_LAST_DIM,
"FUSED_LINEAR" => FusedOps::FUSED_LINEAR,
"RMS_NORM_LAST_DIM" => FusedOps::RMS_NORM_LAST_DIM,
"LAYER_NORM_LAST_DIM" => FusedOps::LAYER_NORM_LAST_DIM,
"ROPE" => FusedOps::ROPE,
"CONV2D" => FusedOps::CONV2D,
"SOFTMAX_LAST_DIM_BACKWARD" => FusedOps::SOFTMAX_LAST_DIM_BACKWARD,
"LAYER_NORM_LAST_DIM_BACKWARD" => FusedOps::LAYER_NORM_LAST_DIM_BACKWARD,
"RMS_NORM_LAST_DIM_BACKWARD" => FusedOps::RMS_NORM_LAST_DIM_BACKWARD,
"REDUCE_MAX_TO_BACKWARD" => FusedOps::REDUCE_MAX_TO_BACKWARD,
"CONV_TRANSPOSE2D" => FusedOps::CONV_TRANSPOSE2D,
"FLASH_ATTN" => FusedOps::FLASH_ATTN,
"PAGED_ATTN" => FusedOps::PAGED_ATTN,
"QMATMUL" => FusedOps::QMATMUL,
"POWI_BACKWARD" => FusedOps::POWI_BACKWARD,
"INPLACE_AFFINE" => FusedOps::INPLACE_AFFINE,
"FUSED_SOFTMAX_CROSS_ENTROPY" => FusedOps::FUSED_SOFTMAX_CROSS_ENTROPY,
"CAUSAL_CONV1D" => FusedOps::CAUSAL_CONV1D,
"SELECTIVE_SCAN" => FusedOps::SELECTIVE_SCAN,
"SSD_CHUNK_SCAN" => FusedOps::SSD_CHUNK_SCAN,
"NF4_MATMUL" => FusedOps::NF4_MATMUL,
"FLASH_ATTN_BACKWARD_Q" => FusedOps::FLASH_ATTN_BACKWARD_Q,
"FLASH_ATTN_BACKWARD_K" => FusedOps::FLASH_ATTN_BACKWARD_K,
"FLASH_ATTN_BACKWARD_V" => FusedOps::FLASH_ATTN_BACKWARD_V,
"AWQ_MATMUL" => FusedOps::AWQ_MATMUL,
_ => return None,
};
Some(id)
}
pub(crate) fn lower_dtype(token: &str, section: &str, operand: &str) -> Result<DType, FkcError> {
let mapped = match token {
"U8" => Some(DType::U8),
"I8" => Some(DType::I8),
"U32" => Some(DType::U32),
"I16" => Some(DType::I16),
"I32" => Some(DType::I32),
"I64" => Some(DType::I64),
"BF16" => Some(DType::BF16),
"F16" => Some(DType::F16),
"F32" => Some(DType::F32),
"F64" => Some(DType::F64),
"F8E4M3" => Some(DType::F8E4M3),
"F6E2M3" => Some(DType::F6E2M3),
"F6E3M2" => Some(DType::F6E3M2),
"F4" => Some(DType::F4),
"F8E8M0" => Some(DType::F8E8M0),
"F8E5M2" => Some(DType::F8E5M2),
"F8E6M2" => Some(DType::F8E6M2),
"BOOL" => Some(DType::Bool),
_ => None,
};
if let Some(dt) = mapped {
match dt {
DType::U8
| DType::I8
| DType::U32
| DType::I16
| DType::I32
| DType::I64
| DType::BF16
| DType::F16
| DType::F32
| DType::F64
| DType::F8E4M3
| DType::F8E5M2
| DType::F6E2M3
| DType::F6E3M2
| DType::F4
| DType::F8E8M0
| DType::F8E6M2
| DType::Bool => (),
};
Ok(dt)
} else if fuel_ir::is_reserved_dtype_token(token) {
Err(FkcError::ReservedScalarType {
section: section.to_string(),
operand: operand.to_string(),
token: token.to_string(),
})
} else {
Err(FkcError::BadScalarType {
section: section.to_string(),
operand: operand.to_string(),
token: token.to_string(),
})
}
}
pub(crate) fn expand_dtype_class(
class: &str,
enumerated: &[DType],
section: &str,
operand: &str,
) -> Result<Vec<DType>, FkcError> {
match class {
"float" => Ok(vec![DType::BF16, DType::F16, DType::F32, DType::F64]),
"int" => Ok(vec![DType::I8, DType::I16, DType::I32, DType::I64]),
"uint" => Ok(vec![DType::U8, DType::U32]),
"any" => Ok(enumerated.to_vec()),
other => Err(FkcError::BadScalarType {
section: section.to_string(),
operand: operand.to_string(),
token: format!("dtype_class={other}"),
}),
}
}
fn lower_backend(s: &str, section: &str) -> Result<BackendId, FkcError> {
let mapped = match s {
"Cpu" => Some(BackendId::Cpu),
"Cuda" => Some(BackendId::Cuda),
"Vulkan" => Some(BackendId::Vulkan),
"Metal" => Some(BackendId::Metal),
_ => None,
};
match mapped {
Some(b) => Ok(b),
None => Err(FkcError::UnknownBackend {
section: section.to_string(),
backend: s.to_string(),
}),
}
}
#[derive(Debug, Clone)]
pub(crate) struct DtypeVariant {
pub(crate) dtypes: KernelDTypes,
pub(crate) fan_dtype: Option<DType>,
}
fn dtype_suffix(dt: DType) -> &'static str {
dt.as_str()
}
struct InputOperand {
name: String,
dtypes: Vec<DType>,
optional: bool,
}
fn resolve_operand_dtypes(
d: &crate::fkc::schema::TensorDesc,
section: &str,
operand: &str,
) -> Result<Vec<DType>, FkcError> {
let mut enumerated: Vec<DType> = Vec::new();
for tok in &d.dtypes {
enumerated.push(lower_dtype(tok, section, operand)?);
}
let resolved = if enumerated.is_empty() {
if let Some(class) = &d.dtype_class {
expand_dtype_class(class, &enumerated, section, operand)?
} else {
enumerated
}
} else {
enumerated
};
if resolved.is_empty() {
return Err(FkcError::BadScalarType {
section: section.to_string(),
operand: operand.to_string(),
token: "<no dtypes and no dtype_class>".to_string(),
});
}
Ok(resolved)
}
fn assemble_dtype_variants(
kernel: &FkcKernel,
section: &str,
) -> Result<(Vec<DtypeVariant>, Vec<ResolvedLayout>), FkcError> {
let mut inputs: Vec<InputOperand> = Vec::new();
let mut layouts: Vec<ResolvedLayout> = Vec::new();
if let Some(accept) = &kernel.accept {
for d in &accept.inputs {
let operand = d.name.as_deref().unwrap_or("<input>");
let resolved = resolve_operand_dtypes(d, section, operand)?;
inputs.push(InputOperand {
name: operand.to_string(),
dtypes: resolved,
optional: d.optional,
});
layouts.push(caps_map::resolve_layout(
d.layout.as_ref(),
section,
operand,
)?);
}
}
let optional_last = inputs.last().map(|op| op.optional).unwrap_or(false);
for (i, op) in inputs.iter().enumerate() {
if op.optional && i != inputs.len() - 1 {
return Err(FkcError::OptionalOperandNotLast {
section: section.to_string(),
operand: op.name.clone(),
});
}
}
let mut fan_set: Option<&[DType]> = None;
for operand in &inputs {
if operand.dtypes.len() > 1 {
match fan_set {
None => fan_set = Some(&operand.dtypes),
Some(existing) => {
if existing != operand.dtypes.as_slice() {
return Err(FkcError::FanoutDtypeMismatch {
section: section.to_string(),
operand: operand.name.clone(),
expected: format!("{existing:?}"),
found: format!("{:?}", operand.dtypes),
});
}
}
}
}
}
let fan_dtypes: Vec<Option<DType>> = match fan_set {
None => vec![None],
Some(set) => set.iter().map(|dt| Some(*dt)).collect(),
};
let per_variant = if optional_last { 2 } else { 1 };
let mut variants: Vec<DtypeVariant> = Vec::with_capacity(fan_dtypes.len() * per_variant);
for fan in fan_dtypes {
let mut input_dtypes: KernelDTypes = SmallVec::new();
for operand in &inputs {
let dt = if operand.dtypes.len() > 1 {
fan.expect("a varying operand implies a fanned dtype")
} else {
operand.dtypes[0]
};
input_dtypes.push(dt);
}
let mut output_dtypes: KernelDTypes = SmallVec::new();
if let Some(ret) = &kernel.return_ {
for out in &ret.outputs {
let operand = out.name.as_deref().unwrap_or("<output>");
if let Some(rule) = out.dtype_rule.as_deref()
&& let Some(dt) = resolve_output_slot_dtype(
rule,
operand,
&inputs,
fan,
optional_last,
section,
)?
{
output_dtypes.push(dt);
}
}
if let Some(bundle) = &ret.bundle
&& let Some((slot, rule)) = bundle_primary_dtype_rule(bundle)
&& let Some(dt) =
resolve_output_slot_dtype(&rule, &slot, &inputs, fan, optional_last, section)?
{
output_dtypes.push(dt);
}
}
if optional_last {
let mut present: KernelDTypes = input_dtypes.clone();
present.extend_from_slice(&output_dtypes);
variants.push(DtypeVariant {
dtypes: present,
fan_dtype: fan,
});
let mut absent: KernelDTypes = SmallVec::new();
absent.extend_from_slice(&input_dtypes[..input_dtypes.len() - 1]);
absent.extend_from_slice(&output_dtypes);
variants.push(DtypeVariant {
dtypes: absent,
fan_dtype: fan,
});
} else {
let mut dtypes: KernelDTypes = input_dtypes;
dtypes.extend_from_slice(&output_dtypes);
variants.push(DtypeVariant {
dtypes,
fan_dtype: fan,
});
}
}
Ok((variants, layouts))
}
fn passthrough_dtype(inputs: &[InputOperand], role: &str, fan: Option<DType>) -> Option<DType> {
let operand_dtype = |op: &InputOperand| -> DType {
if op.dtypes.len() > 1 {
fan.expect("a varying operand implies a fanned dtype")
} else {
op.dtypes[0]
}
};
if let Some(op) = inputs.iter().find(|op| op.name == role) {
Some(operand_dtype(op))
} else {
inputs.first().map(operand_dtype)
}
}
fn resolve_output_slot_dtype(
rule: &str,
operand: &str,
inputs: &[InputOperand],
fan: Option<DType>,
optional_last: bool,
section: &str,
) -> Result<Option<DType>, FkcError> {
match parse_dtype_rule(rule, section, operand)? {
DtypeRule::Fixed(dt) => Ok(Some(dt)),
DtypeRule::Passthrough(role) => {
if optional_last && inputs.last().map(|o| o.name == role).unwrap_or(false) {
return Err(FkcError::PassthroughNamesOptionalOperand {
section: section.to_string(),
role,
});
}
Ok(passthrough_dtype(inputs, &role, fan))
}
DtypeRule::Other => Ok(None),
}
}
fn bundle_primary_dtype_rule(bundle: &serde_yaml_ng::Value) -> Option<(String, String)> {
let serde_yaml_ng::Value::Sequence(slots) = bundle else {
return None;
};
let serde_yaml_ng::Value::Mapping(first) = slots.first()? else {
return None;
};
let rule = first
.get(serde_yaml_ng::Value::String("dtype_rule".into()))
.and_then(|v| v.as_str())?
.to_string();
let name = first
.get(serde_yaml_ng::Value::String("name".into()))
.and_then(|v| v.as_str())
.unwrap_or("<bundle>")
.to_string();
Some((name, rule))
}
enum DtypeRule {
Fixed(DType),
Passthrough(String),
Other,
}
fn parse_dtype_rule(rule: &str, section: &str, operand: &str) -> Result<DtypeRule, FkcError> {
let rule = rule.trim();
if let Some(inner) = rule
.strip_prefix("fixed(")
.and_then(|s| s.strip_suffix(")"))
{
Ok(DtypeRule::Fixed(lower_dtype(
inner.trim(),
section,
operand,
)?))
} else if let Some(inner) = rule
.strip_prefix("passthrough(")
.and_then(|s| s.strip_suffix(")"))
{
Ok(DtypeRule::Passthrough(inner.trim().to_string()))
} else {
Ok(DtypeRule::Other)
}
}
fn compile_cost(block: Option<&CostBlock>, section: &str) -> Result<CompiledCostExpr, FkcError> {
let Some(cost) = block else {
return Ok(CompiledCostExpr::Unknown);
};
let parse = |field: &str, src: Option<&str>| -> Result<CompiledCostExpr, FkcError> {
cost_expr::compile_field(src).map_err(|e| FkcError::CostExprParse {
section: section.to_string(),
field: field.to_string(),
expr: src.unwrap_or("").to_string(),
reason: e.to_string(),
})
};
let flops = parse("flops", cost.flops.as_deref())?;
let _bytes_moved = parse("bytes_moved", cost.bytes_moved.as_deref())?;
if let Some(mem) = &cost.memory
&& let Some(serde_yaml_ng::Value::String(s)) = &mem.device_bytes
{
let _ = parse("memory.device_bytes", Some(s))?;
}
if let Some(serde_yaml_ng::Value::String(s)) = &cost.overhead_ns {
let _ = parse("overhead_ns", Some(s))?;
}
match flops {
CompiledCostExpr::Unknown => Ok(_bytes_moved),
expr => Ok(expr),
}
}
struct Defaults<'a> {
backend: &'a str,
kernel_source: &'a str,
revision_base: &'a str,
}
fn lower_kernel(
kernel: &FkcKernel,
defaults: &Defaults<'_>,
link: &dyn LinkRegistry,
warnings: &mut Vec<crate::fkc::ImportWarning>,
) -> Result<Vec<Resolved>, FkcError> {
let section = kernel.kernel.as_str();
match (kernel.op_kind.as_deref(), kernel.fused_op.as_deref()) {
(Some(op_str), None) => {
let op = lower_op_kind(op_str, section)?;
Ok(lower_primitive(kernel, op, defaults, link)?
.into_iter()
.map(Resolved::Primitive)
.collect())
}
(None, Some(fused_str)) => {
let id = lower_fused_op(fused_str, section)?;
Ok(lower_fused(kernel, id, defaults, link, warnings)?
.into_iter()
.map(Resolved::Fused)
.collect())
}
(op, fused) => Err(FkcError::OpTargetAmbiguous {
section: section.to_string(),
op_kind: op.map(String::from),
fused_op: fused.map(String::from),
}),
}
}
fn lower_primitive(
kernel: &FkcKernel,
op: OpKind,
defaults: &Defaults<'_>,
link: &dyn LinkRegistry,
) -> Result<Vec<ResolvedPrimitive>, FkcError> {
let section = kernel.kernel.as_str();
let backend_str = kernel.backend.as_deref().unwrap_or(defaults.backend);
let backend = lower_backend(backend_str, section)?;
let (variants, layouts) = assemble_dtype_variants(kernel, section)?;
let caps = caps_map::project_kernel_caps(&layouts);
let precision = precision::lower_precision(kernel.precision.as_ref(), section)?;
let cost = compile_cost(kernel.cost.as_ref(), section)?;
let cost_fn = match kernel.cost.as_ref().and_then(|c| c.cost_fn.as_deref()) {
Some(name) => Some(
link.resolve_cost_fn(name)
.ok_or_else(|| FkcError::UnknownCostFn {
section: section.to_string(),
cost_fn: name.to_string(),
})?,
),
None => None,
};
let base_entry_point =
kernel
.entry_point
.as_deref()
.ok_or_else(|| FkcError::MissingRequiredField {
section: section.to_string(),
field: "entry_point".to_string(),
})?;
let kernel_source = kernel
.kernel_source
.as_deref()
.unwrap_or(defaults.kernel_source)
.to_string();
let revision = revhash::compute_revision(kernel, base_entry_point, defaults.revision_base)?;
let mut out = Vec::with_capacity(variants.len());
for variant in variants {
let symbol: std::borrow::Cow<'_, str> = match variant.fan_dtype {
Some(dt) => std::borrow::Cow::Owned(format!("{base_entry_point}_{}", dtype_suffix(dt))),
None => std::borrow::Cow::Borrowed(base_entry_point),
};
let kernel_ref =
link.resolve_primitive(&symbol)
.ok_or_else(|| FkcError::UnknownEntryPoint {
section: section.to_string(),
entry_point: symbol.into_owned(),
})?;
out.push(ResolvedPrimitive {
op,
dtypes: variant.dtypes,
backend,
caps,
layouts: layouts.clone(),
precision,
cost: cost.clone(),
cost_fn,
kernel: kernel_ref,
kernel_source: kernel_source.clone(),
revision,
variant: kernel.variant.clone(),
});
}
Ok(out)
}
fn lower_fused(
kernel: &FkcKernel,
id: FusedOpId,
defaults: &Defaults<'_>,
link: &dyn LinkRegistry,
warnings: &mut Vec<crate::fkc::ImportWarning>,
) -> Result<Vec<ResolvedFused>, FkcError> {
crate::fkc::return_check::cross_check_fused_section(kernel, id, warnings)?;
let section = kernel.kernel.as_str();
let backend_str = kernel.backend.as_deref().unwrap_or(defaults.backend);
let backend = lower_backend(backend_str, section)?;
let (variants, layouts) = assemble_dtype_variants(kernel, section)?;
let caps = caps_map::project_kernel_caps(&layouts);
let precision = precision::lower_precision(kernel.precision.as_ref(), section)?;
let cost = compile_cost(kernel.cost.as_ref(), section)?;
let base_entry_point =
kernel
.entry_point
.as_deref()
.ok_or_else(|| FkcError::MissingRequiredField {
section: section.to_string(),
field: "entry_point".to_string(),
})?;
let kernel_source = kernel
.kernel_source
.as_deref()
.unwrap_or(defaults.kernel_source)
.to_string();
let revision = revhash::compute_revision(kernel, base_entry_point, defaults.revision_base)?;
let mut out = Vec::with_capacity(variants.len());
for variant in variants {
let symbol: std::borrow::Cow<'_, str> = match variant.fan_dtype {
Some(dt) => std::borrow::Cow::Owned(format!("{base_entry_point}_{}", dtype_suffix(dt))),
None => std::borrow::Cow::Borrowed(base_entry_point),
};
let kernel_ref =
link.resolve_fused(&symbol)
.ok_or_else(|| FkcError::UnknownEntryPoint {
section: section.to_string(),
entry_point: symbol.into_owned(),
})?;
out.push(ResolvedFused {
id,
dtypes: variant.dtypes,
backend,
caps,
layouts: layouts.clone(),
precision,
cost: cost.clone(),
kernel: kernel_ref,
kernel_source: kernel_source.clone(),
revision,
variant: kernel.variant.clone(),
bundle_slot_names: crate::fkc::return_check::bundle_slot_names(&kernel.return_),
});
}
Ok(out)
}
pub fn lower_file(
file: &crate::fkc::schema::FkcFile,
link: &dyn LinkRegistry,
warnings: &mut Vec<crate::fkc::ImportWarning>,
) -> Result<Vec<Resolved>, FkcError> {
let provider = &file.front_matter.provider;
let defaults = Defaults {
backend: provider.backend.as_str(),
kernel_source: provider.kernel_source.as_str(),
revision_base: provider.revision_base.as_deref().unwrap_or(""),
};
let mut out = Vec::new();
for kernel in file.kernels.iter().filter(|k| k.registrable) {
out.extend(lower_kernel(kernel, &defaults, link, warnings)?);
}
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::fkc::parse_file;
use std::sync::{Arc, RwLock};
const ELEMENTWISE_BINARY: &str =
include_str!("../../../docs/kernel-contracts/cpu/elementwise-binary.fkc.md");
const QUANT_MATMUL: &str =
include_str!("../../../docs/kernel-contracts/cpu/quant-matmul.fkc.md");
const FUSED_NORM_SOFTMAX: &str =
include_str!("../../../docs/kernel-contracts/fused/norm-softmax.fkc.md");
#[test]
fn a_reserved_dtype_token_declines_distinctly_from_an_unknown_one() {
let reserved = lower_dtype("F8E4M3FNUZ", "sec", "lhs").expect_err("must not lower");
assert!(
matches!(reserved, FkcError::ReservedScalarType { .. }),
"a §6.1-reserved token must decline AS reserved, not as a bad token \
— a contract author told 'bad dtype token' goes hunting for a typo \
that does not exist; got {reserved:?}",
);
let unknown = lower_dtype("NOT_A_DTYPE", "sec", "lhs").expect_err("must not lower");
assert!(
matches!(unknown, FkcError::BadScalarType { .. }),
"an unknown token must stay BadScalarType; got {unknown:?}",
);
}
#[test]
fn every_shared_reserved_token_declines_as_reserved_at_this_site() {
for tok in fuel_ir::RESERVED_DTYPE_TOKENS {
let upper = tok.to_ascii_uppercase();
let err = lower_dtype(&upper, "sec", "lhs")
.expect_err("a reserved token must never lower to a DType");
assert!(
matches!(err, FkcError::ReservedScalarType { .. }),
"{upper} is in RESERVED_DTYPE_TOKENS but declined as {err:?} — this \
site has drifted from the shared list",
);
}
assert!(!fuel_ir::RESERVED_DTYPE_TOKENS.is_empty());
}
#[test]
fn a_reserved_token_is_not_accepted_as_the_token_it_extends() {
assert!(
!matches!(lower_dtype("F8E5M2FNUZ", "sec", "lhs"), Ok(DType::F8E5M2)),
"F8E5M2FNUZ must not lower to its OCP prefix F8E5M2",
);
assert!(
!matches!(lower_dtype("F8E4M3FNUZ", "sec", "lhs"), Ok(DType::F8E4M3)),
"F8E4M3FNUZ must not lower to its OCP prefix F8E4M3",
);
assert_eq!(lower_dtype("F8E5M2", "sec", "lhs").unwrap(), DType::F8E5M2);
assert_eq!(lower_dtype("F8E4M3", "sec", "lhs").unwrap(), DType::F8E4M3);
assert!(
fuel_ir::RESERVED_DTYPE_TOKENS
.iter()
.any(|t| t.eq_ignore_ascii_case("F8E5M2FNUZ")),
);
}
fn dummy_kernel(
_inputs: &[Arc<RwLock<fuel_memory::Storage>>],
_outputs: &mut [Arc<RwLock<fuel_memory::Storage>>],
_layouts: &[fuel_ir::Layout],
_params: &crate::kernel::OpParams,
) -> fuel_ir::Result<()> {
Ok(())
}
struct StubLink;
impl LinkRegistry for StubLink {
fn resolve_primitive(&self, _symbol: &str) -> Option<KernelRef> {
Some(dummy_kernel)
}
fn resolve_fused(&self, _symbol: &str) -> Option<KernelRef> {
Some(dummy_kernel)
}
}
struct EmptyLink;
impl LinkRegistry for EmptyLink {
fn resolve_primitive(&self, _symbol: &str) -> Option<KernelRef> {
None
}
fn resolve_fused(&self, _symbol: &str) -> Option<KernelRef> {
None
}
}
fn lower_all(src: &str, kernel_name: &str) -> Result<Vec<Resolved>, FkcError> {
let file = parse_file(src).expect("parses");
let provider = &file.front_matter.provider;
let defaults = Defaults {
backend: provider.backend.as_str(),
kernel_source: provider.kernel_source.as_str(),
revision_base: provider.revision_base.as_deref().unwrap_or(""),
};
let kernel = file
.kernels
.iter()
.find(|k| k.kernel == kernel_name)
.unwrap_or_else(|| panic!("kernel {kernel_name} present"));
lower_kernel(kernel, &defaults, &StubLink, &mut Vec::new())
}
fn lower_one(src: &str, kernel_name: &str) -> Result<Resolved, FkcError> {
Ok(lower_all(src, kernel_name)?
.into_iter()
.next()
.expect("a section lowers to ≥1 variant"))
}
#[test]
fn lowers_real_add_f32() {
let r = lower_one(ELEMENTWISE_BINARY, "add_f32").expect("add_f32 lowers");
let Resolved::Primitive(p) = r else {
panic!("add_f32 is a primitive");
};
assert_eq!(p.op, OpKind::AddElementwise);
assert_eq!(p.dtypes.as_slice(), &[DType::F32, DType::F32, DType::F32]);
assert_eq!(p.backend, BackendId::Cpu);
assert!(!p.caps.strided_input);
assert!(p.precision.bit_stable_on_same_hardware);
assert_eq!(p.precision.max_ulp, Some(0));
assert_eq!(
p.kernel as *const () as usize,
dummy_kernel as *const () as usize
);
assert!(matches!(p.cost, CompiledCostExpr::Expr(_)));
assert_eq!(p.kernel_source, "portable-cpu");
}
#[test]
fn lowers_binary_chassis_section_fans_out_over_all_dtypes() {
let all = lower_all(ELEMENTWISE_BINARY, "binary").expect("binary lowers");
assert_eq!(all.len(), 4, "4 dtypes ⇒ 4 fanned variants");
let keys: Vec<Vec<DType>> = all
.iter()
.map(|r| {
let Resolved::Primitive(p) = r else {
panic!("binary is a primitive")
};
assert_eq!(p.op, OpKind::AddElementwise);
p.dtypes.to_vec()
})
.collect();
assert_eq!(
keys,
vec![
vec![DType::F32, DType::F32, DType::F32],
vec![DType::F64, DType::F64, DType::F64],
vec![DType::BF16, DType::BF16, DType::BF16],
vec![DType::F16, DType::F16, DType::F16],
],
"each variant keys on [lhs, rhs, out(passthrough lhs)] at that dtype"
);
}
#[test]
fn lowers_real_qmatmul_q4_0() {
let r = lower_one(QUANT_MATMUL, "qmatmul_q4_0_f32").expect("q4_0 lowers");
let Resolved::Primitive(p) = r else {
panic!("QMatMul is a primitive")
};
assert_eq!(p.op, OpKind::QMatMul);
assert_eq!(p.dtypes.as_slice(), &[DType::F32, DType::U8, DType::F32]);
assert_eq!(p.backend, BackendId::Cpu);
assert_eq!(p.cost, CompiledCostExpr::Unknown);
assert!(p.precision.bit_stable_on_same_hardware);
}
#[test]
fn lowers_real_nf4_matmul_f32() {
let r = lower_one(QUANT_MATMUL, "nf4_matmul_f32").expect("nf4 lowers");
let Resolved::Primitive(p) = r else { panic!() };
assert_eq!(p.op, OpKind::Nf4Matmul);
assert_eq!(
p.dtypes.as_slice(),
&[DType::F32, DType::U8, DType::F32, DType::F32]
);
}
#[test]
fn lower_file_lowers_every_kernel() {
let file = parse_file(ELEMENTWISE_BINARY).expect("parses");
let resolved = lower_file(&file, &StubLink, &mut Vec::new()).expect("all lower");
let registrable = file.kernels.iter().filter(|k| k.registrable).count();
assert_eq!(resolved.len(), registrable);
}
#[test]
fn cost_trampoline_evaluates_matmul_declared_flops() {
const MATMUL: &str = include_str!("../../../docs/kernel-contracts/cpu/matmul.fkc.md");
let file = parse_file(MATMUL).expect("matmul contract parses");
let resolved =
lower_file(&file, &StubLink, &mut Vec::new()).expect("matmul contract lowers");
let prim = resolved
.iter()
.find_map(|r| match r {
Resolved::Primitive(p) if p.op == OpKind::MatMul => Some(p),
_ => None,
})
.expect("a MatMul primitive in the matmul contract");
let params = crate::kernel::OpParams::Matmul {
lhs_batch_dims: vec![2],
rhs_batch_dims: vec![2],
m: 3,
n: 4,
k: 5,
m_compute: crate::kernel::MatmulM::All,
};
let est = crate::fkc::cost_estimate(
&prim.cost,
OpKind::MatMul,
&[],
&[DType::F32, DType::F32, DType::F32],
¶ms,
)
.expect("declared matmul cost evaluates");
assert_eq!(est.flops, 240, "2 * batch(2) * m(3) * n(4) * k(5) = 240");
}
#[test]
fn lowers_real_fused_softmax_last_dim() {
let r = lower_one(FUSED_NORM_SOFTMAX, "softmax_last_dim")
.expect("softmax_last_dim (fused) lowers");
let Resolved::Fused(f) = r else {
panic!("softmax_last_dim is a fused op");
};
assert_eq!(f.id, FusedOps::SOFTMAX_LAST_DIM);
assert_eq!(f.backend, BackendId::Cpu);
assert_eq!(f.kernel_source, "portable-cpu");
let file = parse_file(FUSED_NORM_SOFTMAX).expect("parses");
let resolved =
lower_file(&file, &StubLink, &mut Vec::new()).expect("all fused kernels lower");
assert_eq!(
resolved.len(),
file.kernels.len() * 4,
"each fused section fans over its 4 dtypes (8 sections × 4 = 32)",
);
assert!(resolved.iter().all(|r| matches!(r, Resolved::Fused(_))));
}
const FUSED_FANOUT_SYNTH: &str = r#"---
fkc_version: 1
provider:
name: test-provider
backend: Cpu
kernel_source: "test-cpu"
---
# fused dtype-fan synthetic
## fanned_softmax
A blurb.
```fkc
kernel: fanned_softmax
fused_op: SOFTMAX_LAST_DIM
blurb: "synthetic multi-dtype fused softmax"
entry_point: "x::softmax_cpu"
accept:
inputs:
- name: x
dtypes: [F32, F64, BF16, F16]
return:
outputs:
- name: out
dtype_rule: passthrough(x)
```
## single_norm
A blurb.
```fkc
kernel: single_norm
fused_op: RMS_NORM_LAST_DIM
blurb: "synthetic single-dtype fused rms-norm"
entry_point: "x::rms_norm_f32_cpu"
accept:
inputs:
- name: x
dtypes: [F32]
return:
outputs:
- name: out
dtype_rule: passthrough(x)
```
"#;
#[test]
fn fused_multi_dtype_section_fans_into_per_dtype_resolved_fused() {
let all = lower_all(FUSED_FANOUT_SYNTH, "fanned_softmax").expect("fanned_softmax lowers");
assert_eq!(
all.len(),
4,
"4 dtypes ⇒ 4 fanned ResolvedFused (pre-change: 1 representative only)",
);
let keys: Vec<Vec<DType>> = all
.iter()
.map(|r| {
let Resolved::Fused(f) = r else {
panic!("fanned_softmax is a fused op")
};
assert_eq!(f.id, FusedOps::SOFTMAX_LAST_DIM);
assert_eq!(f.backend, BackendId::Cpu);
f.dtypes.to_vec()
})
.collect();
assert_eq!(
keys,
vec![
vec![DType::F32, DType::F32],
vec![DType::F64, DType::F64],
vec![DType::BF16, DType::BF16],
vec![DType::F16, DType::F16],
],
"each variant keys on [x, out(passthrough x)] at that dtype",
);
}
#[test]
fn fused_single_dtype_section_yields_exactly_one_resolved_fused() {
let all = lower_all(FUSED_FANOUT_SYNTH, "single_norm").expect("single_norm lowers");
assert_eq!(
all.len(),
1,
"single-dtype fused section ⇒ exactly one ResolvedFused"
);
let Resolved::Fused(f) = &all[0] else {
panic!("single_norm is a fused op");
};
assert_eq!(f.id, FusedOps::RMS_NORM_LAST_DIM);
assert_eq!(f.dtypes.as_slice(), &[DType::F32, DType::F32]);
}
const PASSTHROUGH_OPTIONAL_SYNTH: &str = r#"---
fkc_version: 1
provider:
name: test-provider
backend: Cpu
kernel_source: "test-cpu"
---
# passthrough-names-optional synthetic
## bad_passthrough
A blurb.
```fkc
kernel: bad_passthrough
op_kind: AddElementwise
blurb: "output passthrough names the optional last operand"
entry_point: "x::add_cpu"
accept:
inputs:
- name: x
dtypes: [F32]
- name: bias
dtypes: [F32]
optional: true
return:
outputs:
- name: out
dtype_rule: passthrough(bias)
```
"#;
#[test]
fn passthrough_naming_the_optional_last_operand_is_rejected() {
let err = lower_all(PASSTHROUGH_OPTIONAL_SYNTH, "bad_passthrough")
.expect_err("passthrough(bias) names the optional last operand");
match err {
FkcError::PassthroughNamesOptionalOperand { ref role, .. } => {
assert_eq!(role, "bias", "the rejected role is the optional operand");
}
other => panic!("expected PassthroughNamesOptionalOperand, got {other:?}"),
}
let ok = PASSTHROUGH_OPTIONAL_SYNTH.replace("passthrough(bias)", "passthrough(x)");
lower_all(&ok, "bad_passthrough")
.expect("passthrough(x) names a REQUIRED operand and must lower cleanly");
}
const BUNDLE_SYNTH: &str = r#"---
fkc_version: 1
provider:
name: test-provider
backend: Cpu
kernel_source: "test-cpu"
---
# bundle key-derivation synthetic
## bundle_op
A blurb.
```fkc
kernel: bundle_op
op_kind: SelectiveScan
blurb: "synthetic multi-output bundle op"
entry_point: "x::bundle_op"
accept:
inputs:
- name: a
dtypes: [F32]
- name: b
dtypes: [F32]
return:
bundle:
- { name: y, dtype_rule: passthrough(a), shape_rule: same_as(a), layout_guarantee: contiguous }
- { name: st, dtype_rule: passthrough(a), shape_rule: same_as(a), layout_guarantee: contiguous }
```
## plain_op
A blurb.
```fkc
kernel: plain_op
op_kind: AddElementwise
blurb: "synthetic no-bundle op"
entry_point: "x::plain_op"
accept:
inputs:
- name: a
dtypes: [F32]
- name: b
dtypes: [F32]
return:
outputs:
- name: out
dtype_rule: passthrough(a)
```
"#;
#[test]
fn return_bundle_appends_primary_output_slot_to_key() {
let r = lower_one(BUNDLE_SYNTH, "bundle_op").expect("bundle_op lowers");
let Resolved::Primitive(p) = r else {
panic!("bundle_op is a primitive");
};
assert_eq!(
p.dtypes.as_slice(),
&[DType::F32, DType::F32, DType::F32],
"a `return.bundle` must append its primary output slot to the key \
(2 inputs + 1 bundle out = 3); the importer previously built the \
2-slot short key `[F32, F32]` from `return.outputs` only",
);
}
#[test]
fn no_bundle_section_key_is_unchanged() {
let r = lower_one(BUNDLE_SYNTH, "plain_op").expect("plain_op lowers");
let Resolved::Primitive(p) = r else {
panic!("plain_op is a primitive");
};
assert_eq!(
p.dtypes.as_slice(),
&[DType::F32, DType::F32, DType::F32],
"a no-bundle section keys on inputs + regular outputs only",
);
}
const VARIANT_SYNTH: &str = r#"---
fkc_version: 1
provider:
name: test-provider
backend: Cpu
kernel_source: "test-cpu"
---
# variant retention synthetic
## tagged_op
A blurb.
```fkc
kernel: tagged_op
op_kind: MatMul
blurb: "synthetic split-K partial variant"
entry_point: "x::tagged_op"
variant: "splitk_partial"
accept:
inputs:
- name: lhs
dtypes: [F32]
- name: rhs
dtypes: [F32]
return:
outputs:
- name: out
dtype_rule: fixed(F32)
```
## untagged_op
A blurb.
```fkc
kernel: untagged_op
op_kind: MatMul
blurb: "synthetic no-variant op"
entry_point: "x::untagged_op"
accept:
inputs:
- name: lhs
dtypes: [F32]
- name: rhs
dtypes: [F32]
return:
outputs:
- name: out
dtype_rule: fixed(F32)
```
"#;
#[test]
fn variant_tag_is_retained_on_resolved_primitive() {
let r = lower_one(VARIANT_SYNTH, "tagged_op").expect("tagged_op lowers");
let Resolved::Primitive(p) = r else {
panic!("tagged_op is a primitive");
};
assert_eq!(
p.variant.as_deref(),
Some("splitk_partial"),
"the opaque `variant:` tag must survive lowering onto \
ResolvedPrimitive.variant (it was dropped)",
);
}
#[test]
fn no_variant_section_retains_none() {
let r = lower_one(VARIANT_SYNTH, "untagged_op").expect("untagged_op lowers");
let Resolved::Primitive(p) = r else {
panic!("untagged_op is a primitive");
};
assert_eq!(p.variant, None, "a section without `variant:` retains None");
}
#[test]
fn bogus_op_kind_is_unknown_op_kind() {
let err = lower_op_kind("NotARealOp", "k").expect_err("bogus op");
assert!(matches!(err, FkcError::UnknownOpKind { .. }), "got {err:?}");
}
#[test]
fn bogus_dtype_is_bad_scalar_type() {
let err = lower_dtype("F99", "k", "lhs").expect_err("bogus dtype");
assert!(matches!(err, FkcError::BadScalarType { .. }), "got {err:?}");
}
#[test]
fn dtype_suffix_is_the_inverse_of_lower_dtype() {
for dt in [
DType::U8,
DType::I8,
DType::U32,
DType::I16,
DType::I32,
DType::I64,
DType::BF16,
DType::F16,
DType::F32,
DType::F64,
DType::F8E4M3,
DType::F6E2M3,
DType::F6E3M2,
DType::F4,
DType::F8E8M0,
DType::F8E6M2,
] {
let suffix = super::dtype_suffix(dt);
let token = suffix.to_uppercase();
assert_eq!(
super::lower_dtype(&token, "k", "op").expect("round-trips"),
dt,
"suffix {suffix:?} (token {token:?}) must lower back to {dt:?}",
);
}
}
#[test]
fn bogus_fused_op_is_unknown_fused_op() {
let err = lower_fused_op("NotAFusedOp", "k").expect_err("bogus fused");
assert!(
matches!(err, FkcError::UnknownFusedOp { .. }),
"got {err:?}"
);
}
#[test]
fn screaming_snake_const_name_resolves_not_pascalcase() {
assert_eq!(
lower_fused_op("SOFTMAX_LAST_DIM", "k").unwrap(),
FusedOps::SOFTMAX_LAST_DIM
);
assert_eq!(
lower_fused_op("FLASH_ATTN", "k").unwrap(),
FusedOps::FLASH_ATTN
);
assert_eq!(lower_fused_op("QMATMUL", "k").unwrap(), FusedOps::QMATMUL);
assert!(lower_fused_op("SoftmaxLastDim", "k").is_err());
assert!(lower_fused_op("FlashAttn", "k").is_err());
}
#[test]
fn every_registry_id_is_reachable_through_table() {
use std::collections::HashSet;
let reachable: HashSet<FusedOpId> = [
"SOFTMAX_LAST_DIM",
"FUSED_LINEAR",
"RMS_NORM_LAST_DIM",
"LAYER_NORM_LAST_DIM",
"ROPE",
"CONV2D",
"SOFTMAX_LAST_DIM_BACKWARD",
"LAYER_NORM_LAST_DIM_BACKWARD",
"RMS_NORM_LAST_DIM_BACKWARD",
"REDUCE_MAX_TO_BACKWARD",
"CONV_TRANSPOSE2D",
"FLASH_ATTN",
"PAGED_ATTN",
"QMATMUL",
"POWI_BACKWARD",
"INPLACE_AFFINE",
"FUSED_SOFTMAX_CROSS_ENTROPY",
"CAUSAL_CONV1D",
"SELECTIVE_SCAN",
"SSD_CHUNK_SCAN",
"NF4_MATMUL",
"FLASH_ATTN_BACKWARD_Q",
"FLASH_ATTN_BACKWARD_K",
"FLASH_ATTN_BACKWARD_V",
"AWQ_MATMUL",
]
.iter()
.map(|n| super::fused_op_id_for_const_name(n).expect("table arm exists"))
.collect();
for entry in fuel_graph::registry::default_registry().entries_iter() {
assert!(
reachable.contains(&entry.id),
"FusedOpId {:?} (registry name {:?}) is registered but NOT reachable through the \
SCREAMING_SNAKE constant-name table — add a `fused_op_id_for_const_name` arm",
entry.id,
entry.name,
);
}
}
#[test]
fn bogus_backend_is_unknown_backend() {
let err = lower_backend("Tpu", "k").expect_err("bogus backend");
assert!(
matches!(err, FkcError::UnknownBackend { .. }),
"got {err:?}"
);
}
#[test]
fn unknown_entry_point_errors() {
let file = parse_file(ELEMENTWISE_BINARY).expect("parses");
let provider = &file.front_matter.provider;
let defaults = Defaults {
backend: provider.backend.as_str(),
kernel_source: provider.kernel_source.as_str(),
revision_base: provider.revision_base.as_deref().unwrap_or(""),
};
let kernel = file.kernels.iter().find(|k| k.kernel == "add_f32").unwrap();
let err = lower_kernel(kernel, &defaults, &EmptyLink, &mut Vec::new())
.expect_err("no entry point");
assert!(
matches!(err, FkcError::UnknownEntryPoint { .. }),
"got {err:?}"
);
}
#[test]
fn malformed_cost_expr_is_cost_expr_parse() {
let mut block = CostBlock {
provenance: Some("declared".into()),
class: None,
cost_fn: None,
flops: Some("2 * * n".into()), bytes_moved: None,
overhead_ns: None,
memory: None,
};
let err = compile_cost(Some(&block), "k").expect_err("malformed flops");
assert!(matches!(err, FkcError::CostExprParse { .. }), "got {err:?}");
block.flops = Some("n".into());
block.bytes_moved = Some("3 * n *".into());
let err = compile_cost(Some(&block), "k").expect_err("malformed bytes_moved");
assert!(matches!(err, FkcError::CostExprParse { field, .. } if field == "bytes_moved"));
}
#[test]
fn both_op_kind_and_fused_op_is_ambiguous() {
let file = parse_file(ELEMENTWISE_BINARY).expect("parses");
let mut kernel = file.kernels[0].clone();
kernel.fused_op = Some("SoftmaxLastDim".into()); let defaults = Defaults {
backend: "Cpu",
kernel_source: "x",
revision_base: "",
};
let err =
lower_kernel(&kernel, &defaults, &StubLink, &mut Vec::new()).expect_err("ambiguous");
assert!(
matches!(err, FkcError::OpTargetAmbiguous { .. }),
"got {err:?}"
);
}
}