progress

lambdaclass · Feb 12, 2024 · 9f9f379 · 9f9f379
1 parent 3ac546e
commit 9f9f379
Showing 1 changed file with 18 additions and 4 deletions.
diff --git a/crates/concrete_codegen_mlir/src/codegen.rs b/crates/concrete_codegen_mlir/src/codegen.rs
@@ -10,7 +10,7 @@ use concrete_session::Session;
 use melior::{
     dialect::{
         arith::{self, CmpiPredicate},
-        cf, func, memref,
+        cf, func, llvm, memref,
     },
     ir::{
         attribute::{
@@ -203,7 +203,7 @@ pub fn compile_function(ctx: FunctionCodegenCtx) -> Result<(), Box<dyn std::erro
         for statement in &block.statements {
             match &statement.kind {
                 concrete_ir::StatementKind::Assign(place, rvalue) => {
-                    let (value, ty) = compile_rvalue(&ctx, mlir_block, rvalue, &locals);
+                    let (value, _ty) = compile_rvalue(&ctx, mlir_block, rvalue, &locals);
                     compile_store_place(
                         &ctx,
                         mlir_block,
@@ -225,8 +225,22 @@ pub fn compile_function(ctx: FunctionCodegenCtx) -> Result<(), Box<dyn std::erro
                     Location::unknown(ctx.context()),
                 ));
             }
-            concrete_ir::TerminatorKind::Return => {}
-            concrete_ir::TerminatorKind::Unreachable => todo!(),
+            concrete_ir::TerminatorKind::Return => {
+                if let Some(ret_local) = ret_local {
+                    let ptr = locals.get(&ret_local).unwrap();
+                    let value = mlir_block
+                        .append_operation(memref::load(*ptr, &[], Location::unknown(ctx.context())))
+                        .result(0)?
+                        .into();
+                    mlir_block.append_operation(func::r#return(
+                        &[value],
+                        Location::unknown(ctx.context()),
+                    ));
+                }
+            }
+            concrete_ir::TerminatorKind::Unreachable => {
+                mlir_block.append_operation(llvm::unreachable(Location::unknown(ctx.context())));
+            }
             concrete_ir::TerminatorKind::Call {
                 func,
                 args,