pliron_llvm/
function_call_utils.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) The pliron contributors
3
4//! Helper functions to call common simple C functions
5
6use pliron::{
7    arg_err,
8    builtin::{
9        op_interfaces::{OneResultInterface, SymbolTableInterface},
10        types::{IntegerType, Signedness},
11    },
12    context::Context,
13    identifier::Identifier,
14    irbuild::inserter::Inserter,
15    result::Result,
16    symbol_table::SymbolTableCollection,
17    r#type::TypeHandle,
18    value::Value,
19};
20
21use crate::{
22    op_interfaces::CastOpInterface,
23    ops::{FuncOp, GepIndex, GetElementPtrOp, PtrToIntOp, ZeroOp},
24    types::{FuncType, PointerType, VoidType},
25};
26
27#[derive(Debug, thiserror::Error)]
28pub enum LookupOrInsertFunctionError {
29    #[error("Symbol '{0}' found but is not a function")]
30    SymbolNotFunction(Identifier),
31    #[error("Existing function '{0}' has a different type than the one being inserted")]
32    FunctionTypeMismatch(Identifier),
33}
34
35/// Looks up a function by name in the given symbol table.
36/// If it exists, checks that its type matches the provided type.
37/// If it doesn't exist, inserts a new function with the given name and type.
38pub fn lookup_or_insert_function(
39    ctx: &mut Context,
40    symbol_table_collection: &mut SymbolTableCollection,
41    symbol_table_op: Box<dyn SymbolTableInterface>,
42    name: Identifier,
43    return_type: TypeHandle,
44    param_types: Vec<TypeHandle>,
45    is_var_arg: bool,
46) -> Result<FuncOp> {
47    let loc = symbol_table_op.loc(ctx);
48    let func_ty = FuncType::get(ctx, return_type, param_types, is_var_arg);
49    let symbol_table = symbol_table_collection.get_symbol_table(ctx, symbol_table_op.clone());
50    if let Some(func) = symbol_table.lookup(&name) {
51        if let Some(func_op) = func.as_any().downcast_ref::<FuncOp>() {
52            let existing_func_ty = func_op.get_type(ctx);
53            if existing_func_ty != func_ty {
54                return arg_err!(loc, LookupOrInsertFunctionError::FunctionTypeMismatch(name));
55            }
56            Ok(*func_op)
57        } else {
58            arg_err!(loc, LookupOrInsertFunctionError::SymbolNotFunction(name))
59        }
60    } else {
61        let func = FuncOp::new(ctx, name.clone(), func_ty);
62        symbol_table.insert(ctx, Box::new(func), None)?;
63        Ok(func)
64    }
65}
66
67/// Get the type used to represet size
68pub fn get_size_type(ctx: &mut Context) -> TypeHandle {
69    IntegerType::get(ctx, 64, Signedness::Signless).into()
70}
71
72/// Get a declaration to the `malloc` function,
73/// inserting it if it doesn't already exist in the symbol table.
74pub fn lookup_or_create_malloc_fn(
75    ctx: &mut Context,
76    symbol_table_collection: &mut SymbolTableCollection,
77    symbol_table_op: Box<dyn SymbolTableInterface>,
78) -> Result<FuncOp> {
79    let size_ty = get_size_type(ctx);
80    let ret_ty = PointerType::get(ctx, 0).into();
81    lookup_or_insert_function(
82        ctx,
83        symbol_table_collection,
84        symbol_table_op,
85        "malloc".try_into().unwrap(),
86        ret_ty,
87        vec![size_ty],
88        false,
89    )
90}
91
92/// Get a declaration to the `free` function,
93/// inserting it if it doesn't already exist in the symbol table.
94pub fn lookup_or_create_free_fn(
95    ctx: &mut Context,
96    symbol_table_collection: &mut SymbolTableCollection,
97    symbol_table_op: Box<dyn SymbolTableInterface>,
98) -> Result<FuncOp> {
99    let ptr_ty = PointerType::get(ctx, 0).into();
100    lookup_or_insert_function(
101        ctx,
102        symbol_table_collection,
103        symbol_table_op,
104        "free".try_into().unwrap(),
105        VoidType::get(ctx).into(),
106        vec![ptr_ty],
107        false,
108    )
109}
110
111/// Compute size of a type in bytes
112pub fn compute_type_size_in_bytes(
113    ctx: &mut Context,
114    inserter: &mut dyn Inserter,
115    ty: TypeHandle,
116) -> Value {
117    // This is LLVM's expansion for sizeof
118    // (as per a comment in MLIR's `ConvertToLLVMPattern::getSizeInBytes`)
119    //   %0 = getelementptr %ty* null, %sizeType 1
120    //   %1 = ptrtoint %ty* %0 to %sizeType
121    let size_ty = get_size_type(ctx);
122    let pointer_ty = PointerType::get(ctx, 0).into();
123    let zero_op = ZeroOp::new(ctx, pointer_ty);
124    inserter.append_op(ctx, &zero_op);
125    let gep_op = GetElementPtrOp::new(
126        ctx,
127        zero_op.get_result(ctx),
128        vec![GepIndex::Constant(1)],
129        ty,
130    );
131    inserter.append_op(ctx, &gep_op);
132    let ptr_to_int_op = PtrToIntOp::new(ctx, gep_op.get_result(ctx), size_ty);
133    inserter.append_op(ctx, &ptr_to_int_op);
134    ptr_to_int_op.get_result(ctx)
135}
136
137#[cfg(all(test, feature = "llvm-sys"))]
138mod tests {
139    use expect_test::expect;
140    use pliron::{
141        builtin::{
142            op_interfaces::{
143                CallOpCallable, OneResultInterface, SingleBlockRegionInterface, SymbolOpInterface,
144            },
145            ops::ModuleOp,
146            types::FP64Type,
147        },
148        context::Context,
149        init_env_logger_for_tests,
150        irbuild::{
151            inserter::{IRInserter, Inserter, OpInsertionPoint},
152            listener::DummyListener,
153        },
154        op::{Op, verify_op},
155        result::ExpectOk,
156    };
157
158    use crate::{
159        function_call_utils::{
160            compute_type_size_in_bytes, get_size_type, lookup_or_create_free_fn,
161            lookup_or_create_malloc_fn,
162        },
163        llvm_sys::{core::LLVMContext, lljit::LLVMLLJIT, target},
164        ops::{CallOp, FuncOp, ReturnOp},
165        to_llvm_ir::convert_module,
166        types::FuncType,
167    };
168
169    #[test]
170    fn test_malloc_and_free_integration() {
171        init_env_logger_for_tests!();
172        let mut ctx = Context::new();
173        let mut symbol_table_collection = pliron::symbol_table::SymbolTableCollection::new();
174
175        // Create a module
176        let module = ModuleOp::new(&mut ctx, "test_module".try_into().unwrap());
177        let module_box = Box::new(module);
178
179        // Get malloc function
180        let malloc_fn =
181            lookup_or_create_malloc_fn(&mut ctx, &mut symbol_table_collection, module_box.clone())
182                .expect("Failed to create malloc function");
183
184        // Get free function
185        let free_fn =
186            lookup_or_create_free_fn(&mut ctx, &mut symbol_table_collection, module_box.clone())
187                .expect("Failed to create free function");
188
189        // Verify both functions were created
190        assert_eq!(
191            malloc_fn.get_symbol_name(&ctx),
192            "malloc".try_into().unwrap()
193        );
194        assert_eq!(free_fn.get_symbol_name(&ctx), "free".try_into().unwrap());
195
196        // Verify calling them again returns the same functions
197        let malloc_fn_2 =
198            lookup_or_create_malloc_fn(&mut ctx, &mut symbol_table_collection, module_box.clone())
199                .expect("Failed to get malloc function again");
200
201        assert!(
202            malloc_fn == malloc_fn_2,
203            "Expected to get the same malloc function on second lookup"
204        );
205
206        // Create a main function
207        let return_type = get_size_type(&mut ctx);
208        let func_ty = FuncType::get(&ctx, return_type, vec![], false);
209        let main_fn = FuncOp::new(&mut ctx, "main".try_into().unwrap(), func_ty);
210        main_fn
211            .get_operation()
212            .insert_at_front(module.get_body(&ctx, 0), &ctx);
213
214        // Insert calls to malloc and free in the entry block of main
215        let entry = main_fn.get_or_create_entry_block(&mut ctx);
216        let mut inserter = IRInserter::<DummyListener>::new(OpInsertionPoint::AtBlockEnd(entry));
217
218        let fp_ty = FP64Type::get(&ctx);
219        let fp_ty_size = compute_type_size_in_bytes(&mut ctx, &mut inserter, fp_ty.into());
220
221        let callee = CallOpCallable::Direct(malloc_fn.get_symbol_name(&ctx));
222        let callee_ty = malloc_fn.get_type(&ctx);
223        let args = vec![fp_ty_size];
224        let malloc_call = CallOp::new(&mut ctx, callee, callee_ty, args);
225        inserter.append_op(&ctx, &malloc_call);
226
227        let ptr_result = malloc_call.get_result(&ctx);
228        let free_callee = CallOpCallable::Direct(free_fn.get_symbol_name(&ctx));
229        let free_callee_ty = free_fn.get_type(&ctx);
230        let free_args = vec![ptr_result];
231        let free_call = CallOp::new(&mut ctx, free_callee, free_callee_ty, free_args);
232        inserter.append_op(&ctx, &free_call);
233
234        let ret_op = ReturnOp::new(&mut ctx, Some(fp_ty_size));
235        inserter.append_op(&ctx, &ret_op);
236
237        verify_op(&module, &ctx).expect_ok(&ctx);
238
239        // Convert to LLVM
240        let llvm_ctx = LLVMContext::default();
241        let llvm_ir = convert_module(&ctx, &llvm_ctx, module).expect_ok(&ctx);
242
243        expect![[r#"
244            ; ModuleID = 'test_module'
245            source_filename = "test_module"
246
247            define i64 @main() {
248            entry_block2v1:
249              %v3 = call ptr @malloc(i64 ptrtoint (ptr getelementptr (double, ptr null, i32 1) to i64))
250              call void @free(ptr %v3)
251              ret i64 ptrtoint (ptr getelementptr (double, ptr null, i32 1) to i64)
252            }
253
254            declare ptr @malloc(i64)
255
256            declare void @free(ptr)
257        "#]]
258        .assert_eq(&llvm_ir.to_string());
259        llvm_ir.verify().expect("Generated LLVM IR is invalid");
260
261        // Execute the LLVM IR using the JIT and check it runs without errors
262        target::initialize_native().expect("Failed to initialize native target for JIT");
263        let jit = LLVMLLJIT::new_with_default_builder().expect("Failed to create LLJIT instance");
264        jit.add_module(llvm_ir)
265            .expect("Failed to add module to JIT");
266        let main_addr = jit
267            .lookup_symbol("main")
268            .expect("Failed to lookup 'main' symbol");
269        let main_fn = unsafe { core::mem::transmute::<u64, fn() -> u64>(main_addr) };
270        let fp_ty_size = main_fn();
271        assert_eq!(
272            fp_ty_size, 8,
273            "Expected size of double type to be 8 bytes, got {}",
274            fp_ty_size
275        );
276    }
277}
pliron_llvm/function_call_utils.rs

pliron_llvm/
function_call_utils.rs
