use super::*;
use crate::meta;
pub fn emit_vulkan_spirv(san_ll: &str) -> Result<Vec<u8>, String> {
let kern = meta::parse_air_kernel_meta(san_ll);
let entry_name = meta::entry_name(san_ll, "kernel");
Ok(emit_vulkan_spirv_with_sidecar(
san_ll,
kern.as_ref(),
entry_name.as_deref(),
kern.as_ref().map(|meta| &meta.buffer_layouts),
)?
.into_bytes())
}
#[cfg(test)]
pub(in crate::native) fn emit_vulkan_spirv_from_typed_blocks(
function_shell: &str,
blocks: Vec<crate::native::cfg::BodyBlock>,
) -> Result<Vec<u8>, String> {
let mut parsed = LlModule::parse(function_shell)?;
let [function] = parsed.functions.as_mut_slice() else {
return Err(format!(
"native emitter: typed-block test adapter needs exactly one function, got {}",
parsed.functions.len()
));
};
function.blocks = blocks;
Ok(finalize_emission(Emitter::new(parsed), None)?.into_bytes())
}
pub(crate) fn emit_vulkan_spirv_with_sidecar(
san_ll: &str,
kern: Option<&meta::KernMeta>,
entry_name: Option<&str>,
buffer_layouts: Option<&HashMap<u32, meta::AirType>>,
) -> Result<crate::emit_sidecar::EmittedSpirv, String> {
emit_vulkan_spirv_inner(san_ll, false, kern, entry_name, buffer_layouts)
}
pub fn emit_vulkan_spirv_construct_tree(san_ll: &str) -> Result<Vec<u8>, String> {
let kern = meta::parse_air_kernel_meta(san_ll);
let entry_name = meta::entry_name(san_ll, "kernel");
Ok(emit_vulkan_spirv_construct_tree_with_sidecar(
san_ll,
kern.as_ref(),
entry_name.as_deref(),
kern.as_ref().map(|meta| &meta.buffer_layouts),
)?
.into_bytes())
}
pub(crate) fn emit_vulkan_spirv_construct_tree_with_sidecar(
san_ll: &str,
kern: Option<&meta::KernMeta>,
entry_name: Option<&str>,
buffer_layouts: Option<&HashMap<u32, meta::AirType>>,
) -> Result<crate::emit_sidecar::EmittedSpirv, String> {
let san_ll = async_copy::lower_simdgroup_async_copy(san_ll);
let san_ll = vec_scalar_merge::lower_vector_scalar_pointer_merge(&san_ll);
let parsed = LlModule::parse_with_stage_meta(&san_ll, kern, entry_name)?;
finalize_emission(Emitter::new(parsed).with_construct_tree(), buffer_layouts)
}
pub fn emit_vulkan_spirv_with_primitive_phi_metadata(san_ll: &str) -> Result<Vec<u8>, String> {
let kern = meta::parse_air_kernel_meta(san_ll);
let entry_name = meta::entry_name(san_ll, "kernel");
Ok(emit_vulkan_spirv_with_primitive_phi_metadata_sidecar(
san_ll,
kern.as_ref(),
entry_name.as_deref(),
kern.as_ref().map(|meta| &meta.buffer_layouts),
)?
.into_bytes())
}
pub(crate) fn emit_vulkan_spirv_with_primitive_phi_metadata_sidecar(
san_ll: &str,
kern: Option<&meta::KernMeta>,
entry_name: Option<&str>,
buffer_layouts: Option<&HashMap<u32, meta::AirType>>,
) -> Result<crate::emit_sidecar::EmittedSpirv, String> {
emit_vulkan_spirv_inner(san_ll, true, kern, entry_name, buffer_layouts)
}
fn emit_vulkan_spirv_inner(
san_ll: &str,
primitive_phi_metadata: bool,
kern: Option<&meta::KernMeta>,
entry_name: Option<&str>,
buffer_layouts: Option<&HashMap<u32, meta::AirType>>,
) -> Result<crate::emit_sidecar::EmittedSpirv, String> {
let san_ll = async_copy::lower_simdgroup_async_copy(san_ll);
let san_ll = vec_scalar_merge::lower_vector_scalar_pointer_merge(&san_ll);
let parsed = if primitive_phi_metadata {
LlModule::parse_with_primitive_phi_metadata_and_stage_meta(&san_ll, kern, entry_name)?
} else {
LlModule::parse_with_stage_meta(&san_ll, kern, entry_name)?
};
finalize_emission(Emitter::new(parsed), buffer_layouts)
}
fn finalize_emission(
emitter: Emitter,
buffer_layouts: Option<&HashMap<u32, meta::AirType>>,
) -> Result<crate::emit_sidecar::EmittedSpirv, String> {
let (mut module, sidecar) = emitter.emit_with_sidecar(buffer_layouts)?;
add_native_module_capabilities(&mut module);
Ok(crate::emit_sidecar::EmittedSpirv { module, sidecar })
}
pub fn emit_vulkan_spirv_inline_sroa(san_ll: &str) -> Result<Vec<u8>, String> {
let kern = meta::parse_air_kernel_meta(san_ll);
let entry_name = meta::entry_name(san_ll, "kernel");
Ok(emit_vulkan_spirv_inline_sroa_with_sidecar(
san_ll,
kern.as_ref(),
entry_name.as_deref(),
kern.as_ref().map(|meta| &meta.buffer_layouts),
)?
.into_bytes())
}
pub(crate) fn emit_vulkan_spirv_inline_sroa_with_sidecar(
san_ll: &str,
kern: Option<&meta::KernMeta>,
entry_name: Option<&str>,
buffer_layouts: Option<&HashMap<u32, meta::AirType>>,
) -> Result<crate::emit_sidecar::EmittedSpirv, String> {
let inlined = inline::inline_nonrecursive_internal_calls(san_ll);
let sroad = sroa::promote_entry_allocas_and_fold_aggregates(&inlined);
let san_ll = vec_scalar_merge::lower_vector_scalar_pointer_merge(&sroad);
let parsed = LlModule::parse_with_stage_meta(&san_ll, kern, entry_name)?;
finalize_emission(Emitter::new(parsed), buffer_layouts)
}
pub fn emit_vulkan_spirv_inline_sroa_raw(san_ll: &str) -> Result<Vec<u8>, String> {
let kern = meta::parse_air_kernel_meta(san_ll);
let entry_name = meta::entry_name(san_ll, "kernel");
Ok(emit_vulkan_spirv_inline_sroa_raw_with_sidecar(
san_ll,
kern.as_ref(),
entry_name.as_deref(),
kern.as_ref().map(|meta| &meta.buffer_layouts),
)?
.into_bytes())
}
pub(crate) fn emit_vulkan_spirv_inline_sroa_raw_with_sidecar(
san_ll: &str,
kern: Option<&meta::KernMeta>,
entry_name: Option<&str>,
buffer_layouts: Option<&HashMap<u32, meta::AirType>>,
) -> Result<crate::emit_sidecar::EmittedSpirv, String> {
let inlined = inline::inline_nonrecursive_internal_calls(san_ll);
let sroad = sroa::promote_entry_allocas_and_fold_aggregates(&inlined);
let san_ll = vec_scalar_merge::lower_vector_scalar_pointer_merge(&sroad);
let mut parsed = LlModule::parse_with_stage_meta(&san_ll, kern, entry_name)?;
mark_all_device_buffers_raw(&mut parsed, false);
finalize_emission(Emitter::new(parsed), buffer_layouts)
}
pub fn emit_vulkan_spirv_inline_sroa_raw_cfg_restructure(san_ll: &str) -> Result<Vec<u8>, String> {
let kern = meta::parse_air_kernel_meta(san_ll);
let entry_name = meta::entry_name(san_ll, "kernel");
Ok(
emit_vulkan_spirv_inline_sroa_raw_cfg_restructure_with_sidecar(
san_ll,
kern.as_ref(),
entry_name.as_deref(),
kern.as_ref().map(|meta| &meta.buffer_layouts),
)?
.into_bytes(),
)
}
pub(crate) fn emit_vulkan_spirv_inline_sroa_raw_cfg_restructure_with_sidecar(
san_ll: &str,
kern: Option<&meta::KernMeta>,
entry_name: Option<&str>,
buffer_layouts: Option<&HashMap<u32, meta::AirType>>,
) -> Result<crate::emit_sidecar::EmittedSpirv, String> {
let inlined = inline::inline_nonrecursive_internal_calls(san_ll);
let sroad = sroa::promote_entry_allocas_and_fold_aggregates(&inlined);
let san_ll = vec_scalar_merge::lower_vector_scalar_pointer_merge(&sroad);
let mut parsed = LlModule::parse_with_stage_meta(&san_ll, kern, entry_name)?;
mark_all_device_buffers_raw(&mut parsed, false);
finalize_emission(Emitter::new(parsed).with_cfg_restructure(), buffer_layouts)
}
pub fn emit_vulkan_spirv_all_buffers_raw(san_ll: &str) -> Result<Vec<u8>, String> {
let kern = meta::parse_air_kernel_meta(san_ll);
let entry_name = meta::entry_name(san_ll, "kernel");
Ok(emit_vulkan_spirv_all_buffers_raw_with_sidecar(
san_ll,
kern.as_ref(),
entry_name.as_deref(),
kern.as_ref().map(|meta| &meta.buffer_layouts),
)?
.into_bytes())
}
pub(crate) fn emit_vulkan_spirv_all_buffers_raw_with_sidecar(
san_ll: &str,
kern: Option<&meta::KernMeta>,
entry_name: Option<&str>,
buffer_layouts: Option<&HashMap<u32, meta::AirType>>,
) -> Result<crate::emit_sidecar::EmittedSpirv, String> {
let san_ll = vec_scalar_merge::lower_vector_scalar_pointer_merge(san_ll);
let mut parsed = LlModule::parse_with_stage_meta(&san_ll, kern, entry_name)?;
mark_all_device_buffers_raw(&mut parsed, false);
finalize_emission(Emitter::new(parsed), buffer_layouts)
}
pub fn emit_vulkan_spirv_all_buffers_raw_with_workgroup(san_ll: &str) -> Result<Vec<u8>, String> {
let kern = meta::parse_air_kernel_meta(san_ll);
let entry_name = meta::entry_name(san_ll, "kernel");
Ok(emit_vulkan_spirv_all_buffers_raw_with_workgroup_sidecar(
san_ll,
kern.as_ref(),
entry_name.as_deref(),
kern.as_ref().map(|meta| &meta.buffer_layouts),
)?
.into_bytes())
}
pub(crate) fn emit_vulkan_spirv_all_buffers_raw_with_workgroup_sidecar(
san_ll: &str,
kern: Option<&meta::KernMeta>,
entry_name: Option<&str>,
buffer_layouts: Option<&HashMap<u32, meta::AirType>>,
) -> Result<crate::emit_sidecar::EmittedSpirv, String> {
let san_ll = vec_scalar_merge::lower_vector_scalar_pointer_merge(san_ll);
let mut parsed = LlModule::parse_with_stage_meta(&san_ll, kern, entry_name)?;
mark_all_device_buffers_raw(&mut parsed, true);
finalize_emission(Emitter::new(parsed), buffer_layouts)
}
pub fn emit_vulkan_spirv_all_buffers_raw_relooper_feed(san_ll: &str) -> Result<Vec<u8>, String> {
let kern = meta::parse_air_kernel_meta(san_ll);
let entry_name = meta::entry_name(san_ll, "kernel");
Ok(
emit_vulkan_spirv_all_buffers_raw_relooper_feed_with_sidecar(
san_ll,
kern.as_ref(),
entry_name.as_deref(),
kern.as_ref().map(|meta| &meta.buffer_layouts),
)?
.into_bytes(),
)
}
pub(crate) fn emit_vulkan_spirv_all_buffers_raw_relooper_feed_with_sidecar(
san_ll: &str,
kern: Option<&meta::KernMeta>,
entry_name: Option<&str>,
buffer_layouts: Option<&HashMap<u32, meta::AirType>>,
) -> Result<crate::emit_sidecar::EmittedSpirv, String> {
let san_ll = vec_scalar_merge::lower_vector_scalar_pointer_merge(san_ll);
let mut parsed = LlModule::parse_with_stage_meta(&san_ll, kern, entry_name)?;
mark_all_device_buffers_raw(&mut parsed, false);
finalize_emission(Emitter::new(parsed).with_relooper_feed(), buffer_layouts)
}
pub fn emit_vulkan_spirv_all_buffers_raw_bda(san_ll: &str) -> Result<Vec<u8>, String> {
let kern = meta::parse_air_kernel_meta(san_ll);
let entry_name = meta::entry_name(san_ll, "kernel");
Ok(emit_vulkan_spirv_all_buffers_raw_bda_with_sidecar(
san_ll,
kern.as_ref(),
entry_name.as_deref(),
kern.as_ref().map(|meta| &meta.buffer_layouts),
)?
.into_bytes())
}
pub(crate) fn emit_vulkan_spirv_all_buffers_raw_bda_with_sidecar(
san_ll: &str,
kern: Option<&meta::KernMeta>,
entry_name: Option<&str>,
buffer_layouts: Option<&HashMap<u32, meta::AirType>>,
) -> Result<crate::emit_sidecar::EmittedSpirv, String> {
let san_ll = vec_scalar_merge::lower_vector_scalar_pointer_merge(san_ll);
let mut parsed = LlModule::parse_with_stage_meta(&san_ll, kern, entry_name)?;
mark_all_device_buffers_raw(&mut parsed, false);
finalize_emission(
Emitter::new(parsed).with_bda_device_pointers(),
buffer_layouts,
)
}
fn mark_all_device_buffers_raw(parsed: &mut LlModule, include_workgroup: bool) {
let keys: Vec<(String, String)> = parsed
.functions
.iter()
.flat_map(|f| {
f.params.iter().filter_map(move |(name, ty)| {
let raw = matches!(ty, ir::LlType::Ptr(1 | 2))
|| (include_workgroup && matches!(ty, ir::LlType::Ptr(3)));
raw.then(|| (f.name.clone(), name.clone()))
})
})
.collect();
for key in keys {
parsed.raw_buffer_params.insert(key);
}
}