use super::judge::{HashMapJudge, JudgeOracle};
use fuel_ir::DType;
use fuel_ir::dispatch::{OpKind, PROFILE_REPORT_VERSION, ProfileReport, SizeClass};
use fuel_ir::probe::BackendId;
#[derive(Debug, Default, Clone)]
pub struct ProfileJudgeOracle {
inner: HashMapJudge,
}
impl ProfileJudgeOracle {
pub fn from_report(report: &ProfileReport) -> Self {
let mut inner = HashMapJudge::new();
if report.version != PROFILE_REPORT_VERSION {
return Self { inner };
}
for e in &report.entries {
debug_assert_eq!(
e.backend, e.device.backend,
"ProfileEntry backend ({:?}) and device.backend ({:?}) disagree",
e.backend, e.device.backend
);
let prev = inner.measured_latency_ns(
e.op,
e.dtype,
e.size_class,
e.backend,
&e.kernel_source,
e.kernel_revision_hash,
);
if prev.is_none_or(|p| e.latency_ns < p) {
inner.insert(
e.op,
e.dtype,
e.size_class,
e.backend,
&e.kernel_source,
e.latency_ns,
e.kernel_revision_hash,
);
}
}
Self { inner }
}
pub fn len(&self) -> usize {
self.inner.len()
}
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
}
impl JudgeOracle for ProfileJudgeOracle {
fn measured_latency_ns(
&self,
op: OpKind,
dtype: DType,
size_class: SizeClass,
backend: BackendId,
kernel_source: &str,
kernel_revision_hash: u64,
) -> Option<u64> {
self.inner.measured_latency_ns(
op,
dtype,
size_class,
backend,
kernel_source,
kernel_revision_hash,
)
}
}
#[cfg(test)]
mod tests {
use super::*;
use fuel_ir::dispatch::ProfileEntry;
fn test_equiv_key(backend: BackendId) -> fuel_ir::probe::EquivalenceKey {
fuel_ir::probe::EquivalenceKey {
backend,
vendor_id: 0,
device_id: 0,
compute_capability: None,
driver_version: String::new(),
}
}
fn entry(
backend: BackendId,
op: OpKind,
size: u32,
device_index: u32,
latency: u64,
kernel_source: &str,
) -> ProfileEntry {
ProfileEntry {
op,
dtype: DType::F32,
size_class: SizeClass(size),
backend,
device_index,
latency_ns: latency,
iterations: 7,
max_rel_error: 1e-6,
kernel_source: kernel_source.to_string(),
kernel_revision_hash: 0,
device: test_equiv_key(backend),
}
}
#[test]
fn empty_report_misses_everywhere() {
let report = ProfileReport {
version: PROFILE_REPORT_VERSION,
entries: vec![],
};
let oracle = ProfileJudgeOracle::from_report(&report);
assert!(oracle.is_empty());
assert!(
oracle
.measured_latency_ns(
OpKind::MatMul,
DType::F32,
SizeClass(16),
BackendId::Cpu,
"",
0,
)
.is_none()
);
}
#[test]
fn sibling_kernel_sources_do_not_collide() {
let report = ProfileReport {
version: PROFILE_REPORT_VERSION,
entries: vec![
entry(
BackendId::Cpu,
OpKind::MatMul,
16,
0,
1_000_000,
"portable-cpu",
),
entry(BackendId::Cpu, OpKind::MatMul, 16, 0, 200_000, "aocl"),
],
};
let oracle = ProfileJudgeOracle::from_report(&report);
assert_eq!(oracle.len(), 2);
let cell = |src: &str| {
oracle.measured_latency_ns(
OpKind::MatMul,
DType::F32,
SizeClass(16),
BackendId::Cpu,
src,
0,
)
};
assert_eq!(cell("portable-cpu"), Some(1_000_000));
assert_eq!(cell("aocl"), Some(200_000));
assert!(cell("mkl").is_none());
assert!(cell("").is_none());
}
#[test]
fn duplicate_keys_keep_minimum_latency() {
let report = ProfileReport {
version: PROFILE_REPORT_VERSION,
entries: vec![
entry(BackendId::Cuda, OpKind::MatMul, 20, 0, 300_000, "cublas"),
entry(BackendId::Cuda, OpKind::MatMul, 20, 1, 900_000, "cublas"),
],
};
let oracle = ProfileJudgeOracle::from_report(&report);
assert_eq!(oracle.len(), 1);
assert_eq!(
oracle.measured_latency_ns(
OpKind::MatMul,
DType::F32,
SizeClass(20),
BackendId::Cuda,
"cublas",
0,
),
Some(300_000),
);
let reversed = ProfileReport {
version: PROFILE_REPORT_VERSION,
entries: vec![
entry(BackendId::Cuda, OpKind::MatMul, 20, 1, 900_000, "cublas"),
entry(BackendId::Cuda, OpKind::MatMul, 20, 0, 300_000, "cublas"),
],
};
let oracle = ProfileJudgeOracle::from_report(&reversed);
assert_eq!(
oracle.measured_latency_ns(
OpKind::MatMul,
DType::F32,
SizeClass(20),
BackendId::Cuda,
"cublas",
0,
),
Some(300_000),
);
}
#[test]
fn miss_on_any_axis_returns_none() {
let report = ProfileReport {
version: PROFILE_REPORT_VERSION,
entries: vec![entry(BackendId::Cpu, OpKind::MatMul, 16, 0, 1_000, "")],
};
let oracle = ProfileJudgeOracle::from_report(&report);
let hit = oracle.measured_latency_ns(
OpKind::MatMul,
DType::F32,
SizeClass(16),
BackendId::Cpu,
"",
0,
);
assert_eq!(hit, Some(1_000));
assert!(
oracle
.measured_latency_ns(
OpKind::MatMul,
DType::F32,
SizeClass(16),
BackendId::Cuda,
"",
0,
)
.is_none()
);
assert!(
oracle
.measured_latency_ns(
OpKind::MatMul,
DType::F64,
SizeClass(16),
BackendId::Cpu,
"",
0,
)
.is_none()
);
assert!(
oracle
.measured_latency_ns(
OpKind::MatMul,
DType::F32,
SizeClass(17),
BackendId::Cpu,
"",
0,
)
.is_none()
);
assert!(
oracle
.measured_latency_ns(
OpKind::AddElementwise,
DType::F32,
SizeClass(16),
BackendId::Cpu,
"",
0,
)
.is_none()
);
}
#[test]
fn version_mismatch_yields_empty_oracle() {
let report = ProfileReport {
version: PROFILE_REPORT_VERSION + 1,
entries: vec![entry(BackendId::Cpu, OpKind::MatMul, 16, 0, 1_000, "")],
};
let oracle = ProfileJudgeOracle::from_report(&report);
assert!(oracle.is_empty());
}
#[test]
fn adapter_attaches_to_plan_options() {
let report = ProfileReport {
version: PROFILE_REPORT_VERSION,
entries: vec![entry(BackendId::Cpu, OpKind::MatMul, 16, 0, 1_000, "aocl")],
};
let oracle = ProfileJudgeOracle::from_report(&report);
let opts = crate::plan::PlanOptions::new().with_judge(&oracle);
let queried = opts
.judge
.expect("with_judge sets the oracle")
.measured_latency_ns(
OpKind::MatMul,
DType::F32,
SizeClass(16),
BackendId::Cpu,
"aocl",
0,
);
assert_eq!(queried, Some(1_000));
}
#[test]
fn sibling_kernel_revisions_do_not_collide_via_report() {
let mk = |rev: u64, latency: u64| ProfileEntry {
op: OpKind::MatMul,
dtype: DType::F32,
size_class: SizeClass(16),
backend: BackendId::Cuda,
device_index: 0,
latency_ns: latency,
iterations: 7,
max_rel_error: 1e-6,
kernel_source: "baracuda".to_string(),
kernel_revision_hash: rev,
device: test_equiv_key(BackendId::Cuda),
};
let report = ProfileReport {
version: PROFILE_REPORT_VERSION,
entries: vec![
mk(0x1111_1111_1111_1111, 5_000),
mk(0x2222_2222_2222_2222, 9_000),
],
};
let oracle = ProfileJudgeOracle::from_report(&report);
assert_eq!(
oracle.len(),
2,
"distinct non-zero revisions are distinct cells"
);
let cell = |rev: u64| {
oracle.measured_latency_ns(
OpKind::MatMul,
DType::F32,
SizeClass(16),
BackendId::Cuda,
"baracuda",
rev,
)
};
assert_eq!(cell(0x1111_1111_1111_1111), Some(5_000));
assert_eq!(cell(0x2222_2222_2222_2222), Some(9_000));
assert!(cell(0).is_none());
}
}