lift_export/

llvm.rs

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);

        // Declare runtime functions
        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 &region.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);

                                        // Emit input argument gathering
                                        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
                                        );
                                        // Annotate the LLVM type of the first
                                        // input/result (informational; the
                                        // runtime calls take pointers).
                                        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()
    }
}