use baracuda_kernels_types::{ElementKind, OperandDesc};
use fuel_graph::registry::FusedOpId;
use fuel_ir::probe::BackendId;
use fuel_ir::{DType, Error, Result};
use fuel_kernel_seam::{JitRequest, JitResponse, SynthArtifact, Synthesizer};
use crate::kernel::KernelRef;
use crate::runtime_fused_kernels::adopt_runtime_fused;
pub(crate) fn element_kind_to_dtype(ek: ElementKind) -> Option<DType> {
Some(match ek {
ElementKind::U8 => DType::U8,
ElementKind::I8 => DType::I8,
ElementKind::I32 => DType::I32,
ElementKind::I64 => DType::I64,
ElementKind::Bf16 => DType::BF16,
ElementKind::F16 => DType::F16,
ElementKind::F32 => DType::F32,
ElementKind::F64 => DType::F64,
ElementKind::U32 => DType::U32,
ElementKind::Fp8E4M3FN => DType::F8E4M3,
ElementKind::Fp8E5M2 => DType::F8E5M2,
ElementKind::Bool => DType::Bool,
ElementKind::F32Strict
| ElementKind::I4
| ElementKind::U4
| ElementKind::B1
| ElementKind::Complex64
| ElementKind::Complex128 => return None,
ElementKind::U16
| ElementKind::U64
| ElementKind::Fp8E4M3FNUZ
| ElementKind::Fp8E5M2FNUZ => {
return None;
}
ElementKind::I16 | ElementKind::F8E8M0 | ElementKind::F8E6M2 => return None,
})
}
pub(crate) fn dtype_to_element_kind(dt: DType) -> Option<ElementKind> {
Some(match dt {
DType::U8 => ElementKind::U8,
DType::I8 => ElementKind::I8,
DType::I32 => ElementKind::I32,
DType::I64 => ElementKind::I64,
DType::BF16 => ElementKind::Bf16,
DType::F16 => ElementKind::F16,
DType::F32 => ElementKind::F32,
DType::F64 => ElementKind::F64,
DType::U32 => ElementKind::U32,
DType::F8E4M3 => ElementKind::Fp8E4M3FN,
DType::F8E5M2 => ElementKind::Fp8E5M2,
DType::Bool => ElementKind::Bool,
DType::I16 | DType::F6E2M3 | DType::F6E3M2 | DType::F4 | DType::F8E8M0 | DType::F8E6M2 => {
return None;
}
})
}
fn operand_dtypes(operands: &[OperandDesc]) -> Option<Vec<DType>> {
operands
.iter()
.map(|o| element_kind_to_dtype(o.dtype))
.collect()
}
#[derive(Clone, Copy, Debug)]
pub struct Adopted {
pub id: FusedOpId,
pub kernel: KernelRef,
}
pub fn adopt_from_response(
synth: &dyn Synthesizer,
req: &JitRequest,
backend: BackendId,
load_kernel: impl FnOnce(&SynthArtifact) -> Result<KernelRef>,
) -> Result<Option<Adopted>> {
let entry_point = match synth.synthesize(req) {
JitResponse::Synthesized { entry_point } => entry_point,
JitResponse::Declined { .. } => return Ok(None),
};
let art = synth.take_kernel(&entry_point).ok_or_else(|| {
Error::Msg(format!(
"take_kernel({entry_point}): synthesizer retained nothing"
))
})?;
let kernel = load_kernel(&art)?;
let Some(dtypes) = operand_dtypes(&req.operands) else {
return Ok(None);
};
Ok(
adopt_runtime_fused(entry_point, req.region.clone(), kernel, dtypes, backend)
.map(|id| Adopted { id, kernel }),
)
}
#[cfg(test)]
mod tests {
#[test]
fn dtype_and_element_kind_round_trip_exactly() {
use super::{dtype_to_element_kind, element_kind_to_dtype};
use fuel_ir::DType;
const MAPPED: &[DType] = &[
DType::U8,
DType::I8,
DType::I32,
DType::I64,
DType::BF16,
DType::F16,
DType::F32,
DType::F64,
DType::U32,
DType::F8E4M3,
DType::F8E5M2,
];
for &dt in MAPPED {
let ek = dtype_to_element_kind(dt)
.unwrap_or_else(|| panic!("{dt:?} must have a seam spelling"));
assert_eq!(
element_kind_to_dtype(ek),
Some(dt),
"{dt:?} -> {ek:?} -> back must be the identity; a non-identity round trip means a request describes operands the caller does not have",
);
}
assert!(
dtype_to_element_kind(DType::I16).is_none(),
"I16 has no ElementKind spelling (the seam has S8/U8/I32/I64/U32, no 16-bit int) — it must decline, not substitute — a lossy mapping asks about a different kernel than the caller's",
);
}
#[test]
fn bool_maps_to_the_seam_bool_and_did_not_stay_declined() {
use super::{dtype_to_element_kind, element_kind_to_dtype};
use baracuda_kernels_types::ElementKind;
use fuel_ir::DType;
assert_eq!(dtype_to_element_kind(DType::Bool), Some(ElementKind::Bool));
assert_eq!(element_kind_to_dtype(ElementKind::Bool), Some(DType::Bool));
}
#[test]
fn fp8_maps_to_baracuda_ocp_element_kinds() {
use super::{dtype_to_element_kind, element_kind_to_dtype};
use baracuda_kernels_types::ElementKind;
use fuel_ir::DType;
assert_eq!(
dtype_to_element_kind(DType::F8E4M3),
Some(ElementKind::Fp8E4M3FN)
);
assert_eq!(
dtype_to_element_kind(DType::F8E5M2),
Some(ElementKind::Fp8E5M2)
);
assert_eq!(
element_kind_to_dtype(ElementKind::Fp8E4M3FN),
Some(DType::F8E4M3)
);
assert_eq!(
element_kind_to_dtype(ElementKind::Fp8E5M2),
Some(DType::F8E5M2)
);
}
use super::*;
use fuel_graph::jit::{OpAttrs, OpTag, PatternNode};
use fuel_kernel_seam::{ArtifactKind, JitBudget, LinkEntry};
use std::sync::{Arc, Mutex, RwLock as StdRwLock};
fn noop_kernel(
_inputs: &[Arc<StdRwLock<fuel_memory::Storage>>],
_outputs: &mut [Arc<StdRwLock<fuel_memory::Storage>>],
_layouts: &[fuel_ir::Layout],
_params: &crate::kernel::OpParams,
) -> Result<()> {
Ok(())
}
fn abs_sub() -> PatternNode {
PatternNode::Op {
op: OpTag::Abs,
attrs: OpAttrs::default(),
operands: vec![PatternNode::Op {
op: OpTag::Sub,
attrs: OpAttrs::default(),
operands: vec![
PatternNode::Bind { index: 0 },
PatternNode::Bind { index: 1 },
],
}],
}
}
fn artifact(entry_point: &str) -> SynthArtifact {
SynthArtifact {
artifact: vec![0xCA, 0xFE],
kind: ArtifactKind::Ptx,
link: LinkEntry {
entry_point: entry_point.into(),
symbol: "k".into(),
structure_key: "elementwise:f32".into(),
revision_hash: 1,
},
contract: "## fused_op\ncost: n\n".into(),
}
}
struct MockSynth {
decline: bool,
art: Mutex<Option<SynthArtifact>>,
}
impl Synthesizer for MockSynth {
fn synthesize(&self, _req: &JitRequest) -> JitResponse {
if self.decline {
JitResponse::Declined {
reason: "mock decline".into(),
}
} else {
JitResponse::Synthesized {
entry_point: "mock::abs_sub".into(),
}
}
}
fn take_kernel(&self, _entry_point: &str) -> Option<SynthArtifact> {
self.art.lock().unwrap().take()
}
}
fn req() -> JitRequest {
JitRequest::new(
abs_sub(),
vec![
OperandDesc::new(1, &[4], &[1], ElementKind::F32, 256),
OperandDesc::new(1, &[4], &[1], ElementKind::F32, 256),
],
baracuda_kernels_types::ArchSku::Sm89,
JitBudget {
max_compile_ms: 250,
},
)
}
#[test]
fn adopts_a_synthesized_kernel_end_to_end() {
let synth = MockSynth {
decline: false,
art: Mutex::new(Some(artifact("mock::abs_sub"))),
};
let adopted = adopt_from_response(&synth, &req(), BackendId::Cpu, |_art| {
Ok(noop_kernel as KernelRef)
})
.expect("no error")
.expect("synthesized ⇒ adopted");
let id = adopted.id;
assert!(id.is_runtime(), "adopted a runtime FusedOpId");
assert!(
std::ptr::fn_addr_eq(adopted.kernel, noop_kernel as KernelRef),
"adoption returned the kernel this call loaded",
);
assert!(
crate::runtime_fused_kernels::fused_kernel_available(id, BackendId::Cpu),
"the adopted op's kernel is now visible to the capability gate",
);
}
#[test]
fn declined_synthesis_adopts_nothing() {
let synth = MockSynth {
decline: true,
art: Mutex::new(None),
};
let out = adopt_from_response(&synth, &req(), BackendId::Cpu, |_art| {
panic!("load_kernel must not run on a decline")
})
.expect("no error");
assert!(out.is_none(), "declined ⇒ no adoption, no kernel load");
}
}