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

pliron_llvm/
function_call_utils.rs
