use lift_core::context::Context;
use lift_core::types::{CoreType, TypeData};
use std::fmt::Write;
use thiserror::Error;
#[derive(Debug, Error)]
pub enum LlvmExportError {
#[error("Unsupported operation for LLVM export: {0}")]
UnsupportedOp(String),
#[error("Export error: {0}")]
General(String),
}
#[derive(Debug)]
pub struct LlvmExporter;
impl LlvmExporter {
pub fn new() -> Self {
Self
}
fn llvm_type_for_value(&self, ctx: &Context, val_key: lift_core::values::ValueKey) -> String {
if let Some(val) = ctx.get_value(val_key) {
match ctx.resolve_type(val.ty) {
CoreType::Opaque {
data: TypeData::Tensor(info),
..
} => {
let mut elems = 1usize;
for dim in &info.shape {
if let Some(s) = dim.static_value() {
elems = elems.saturating_mul(s);
}
}
let dt = match info.dtype.byte_size() {
2 => "half",
4 => "float",
8 => "double",
_ => "float",
};
format!("<{} x {}>", elems, dt)
}
CoreType::Float { bits: 32 } => "float".to_string(),
CoreType::Float { bits: 64 } => "double".to_string(),
CoreType::Integer { bits, .. } => format!("i{}", bits),
CoreType::Boolean => "i1".to_string(),
_ => "ptr".to_string(),
}
} else {
"ptr".to_string()
}
}
fn runtime_call_for_op(&self, op_name: &str) -> &str {
match op_name {
"tensor.matmul" => "lift_rt_matmul",
"tensor.add" => "lift_rt_add",
"tensor.sub" => "lift_rt_sub",
"tensor.mul" => "lift_rt_mul",
"tensor.div" => "lift_rt_div",
"tensor.relu" => "lift_rt_relu",
"tensor.gelu" => "lift_rt_gelu",
"tensor.silu" => "lift_rt_silu",
"tensor.softmax" => "lift_rt_softmax",
"tensor.sigmoid" => "lift_rt_sigmoid",
"tensor.tanh" => "lift_rt_tanh",
"tensor.layernorm" => "lift_rt_layernorm",
"tensor.rmsnorm" => "lift_rt_rmsnorm",
"tensor.batchnorm" => "lift_rt_batchnorm",
"tensor.conv2d" => "lift_rt_conv2d",
"tensor.conv1d" => "lift_rt_conv1d",
"tensor.maxpool2d" => "lift_rt_maxpool2d",
"tensor.avgpool2d" => "lift_rt_avgpool2d",
"tensor.global_avgpool" => "lift_rt_global_avgpool",
"tensor.attention" => "lift_rt_attention",
"tensor.multi_head_attention" => "lift_rt_multi_head_attention",
"tensor.grouped_query_attention" => "lift_rt_grouped_query_attention",
"tensor.flash_attention" => "lift_rt_flash_attention",
"tensor.cross_attention" => "lift_rt_cross_attention",
"tensor.sliding_window_attention" => "lift_rt_sliding_window_attention",
"tensor.paged_attention" => "lift_rt_paged_attention",
"tensor.embedding" => "lift_rt_embedding",
"tensor.linear" => "lift_rt_linear",
"tensor.reshape" => "lift_rt_reshape",
"tensor.transpose" => "lift_rt_transpose",
"tensor.concat" => "lift_rt_concat",
"tensor.moe_dispatch" => "lift_rt_moe_dispatch",
"tensor.moe_combine" => "lift_rt_moe_combine",
"tensor.fused_matmul_bias_relu" => "lift_rt_fused_matmul_bias_relu",
"tensor.fused_matmul_bias" => "lift_rt_fused_matmul_bias",
"tensor.fused_linear_gelu" => "lift_rt_fused_linear_gelu",
"tensor.fused_linear_silu" => "lift_rt_fused_linear_silu",
"tensor.quantize" => "lift_rt_quantize",
"tensor.dequantize" => "lift_rt_dequantize",
_ => "lift_rt_generic_op",
}
}
pub fn export(&self, ctx: &Context) -> Result<String, LlvmExportError> {
let mut output = String::new();
let _ = writeln!(output, "; LIFT IR -> LLVM IR export");
let _ = writeln!(output, "; Generated by LIFT compiler framework");
let _ = writeln!(output, "; Target: GPU runtime (cuBLAS/cuDNN backend)");
let _ = writeln!(output);
let _ = writeln!(output, "; === Runtime function declarations ===");
let rt_funcs = [
"lift_rt_matmul",
"lift_rt_add",
"lift_rt_sub",
"lift_rt_mul",
"lift_rt_div",
"lift_rt_relu",
"lift_rt_gelu",
"lift_rt_silu",
"lift_rt_softmax",
"lift_rt_sigmoid",
"lift_rt_tanh",
"lift_rt_layernorm",
"lift_rt_rmsnorm",
"lift_rt_batchnorm",
"lift_rt_conv2d",
"lift_rt_conv1d",
"lift_rt_maxpool2d",
"lift_rt_avgpool2d",
"lift_rt_global_avgpool",
"lift_rt_attention",
"lift_rt_multi_head_attention",
"lift_rt_grouped_query_attention",
"lift_rt_flash_attention",
"lift_rt_cross_attention",
"lift_rt_sliding_window_attention",
"lift_rt_paged_attention",
"lift_rt_embedding",
"lift_rt_linear",
"lift_rt_reshape",
"lift_rt_transpose",
"lift_rt_concat",
"lift_rt_moe_dispatch",
"lift_rt_moe_combine",
"lift_rt_fused_matmul_bias_relu",
"lift_rt_fused_matmul_bias",
"lift_rt_fused_linear_gelu",
"lift_rt_fused_linear_silu",
"lift_rt_quantize",
"lift_rt_dequantize",
"lift_rt_generic_op",
];
for f in &rt_funcs {
let _ = writeln!(output, "declare ptr @{}(ptr, ptr, ptr, i64, i64)", f);
}
let _ = writeln!(output);
for module in &ctx.modules {
let name = ctx.strings.resolve(module.name);
let _ = writeln!(output, "; === Module: {} ===", name);
let _ = writeln!(output);
for func in &module.functions {
let fname = ctx.strings.resolve(func.name);
let _ = write!(output, "define ptr @{}(", fname);
for (i, _param) in func.params.iter().enumerate() {
if i > 0 {
let _ = write!(output, ", ");
}
let _ = write!(output, "ptr %arg{}", i);
}
let _ = writeln!(output, ") {{");
let _ = writeln!(output, "entry:");
let mut result_counter = 0usize;
if let Some(region_key) = func.body {
if let Some(region) = ctx.get_region(region_key) {
for &block_key in ®ion.blocks {
if let Some(block) = ctx.get_block(block_key) {
for &op_key in &block.ops {
if let Some(op) = ctx.get_op(op_key) {
let op_name = ctx.strings.resolve(op.name).to_string();
let rt_func = self.runtime_call_for_op(&op_name);
let in0 = if !op.inputs.is_empty() {
format!("%arg{}", 0)
} else {
"null".to_string()
};
let in1 = if op.inputs.len() > 1 {
format!("%arg{}", 1)
} else {
"null".to_string()
};
let num_in = op.inputs.len() as i64;
let num_out = op.results.len() as i64;
let _ = writeln!(
output,
" ; {} ({} inputs -> {} outputs)",
op_name, num_in, num_out
);
if let Some(&first_in) = op.inputs.first() {
let ty = self.llvm_type_for_value(ctx, first_in);
let _ = writeln!(output, " ; input type: {}", ty);
}
let _ = writeln!(output,
" %r{} = call ptr @{}(ptr {}, ptr {}, ptr null, i64 {}, i64 {})",
result_counter, rt_func, in0, in1, num_in, num_out);
result_counter += 1;
}
}
}
}
}
}
if result_counter > 0 {
let _ = writeln!(output, " ret ptr %r{}", result_counter - 1);
} else {
let _ = writeln!(output, " ret ptr null");
}
let _ = writeln!(output, "}}");
let _ = writeln!(output);
}
}
Ok(output)
}
}
impl Default for LlvmExporter {
fn default() -> Self {
Self::new()
}
}