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

pliron_llvm/
function_call_utils.rs
