feat: add type checks for MLIR codegen types

Author: lucarlig

examples/error-examples/assign_to_shared_ref.desc

@@ -0,0 +1,7 @@
+fn assign_to_shared_ref<n: nat, r: prv>(
+    a: &r shrd gpu.global [i16; 16],
+    b: &r shrd gpu.global [i16; 16]
+) -[grid: gpu.grid<X<1>, X<16>>]-> () {
+    b = a;  // This should fail - cannot assign to shared reference
+    ()
+}

examples/error-examples/vec_add_memory_issue.desc

@@ -0,0 +1,50 @@
+// Vector addition kernel demonstrating Descend's safe GPU programming model
+// This function showcases extended borrow checking, memory safety, and execution context tracking
+
+// Generic function with type parameters:
+// - n: nat - Natural number parameter (for array size, though not used in this specific function)
+// - r: prv - Provenance parameter tracking memory region/lifetime for all references
+fn add<n: nat, r: prv>(
+    // Shared reference to first input vector - multiple threads can read simultaneously
+    // Memory space: gpu.global (GPU global memory)
+    // Ownership: shrd (shared) - prevents write-after-read data races
+    // Type: 16-element array of 16-bit signed integers
+    a: &r shrd gpu.global [i16; 16],
+    
+    // Shared reference to second input vector - multiple threads can read simultaneously  
+    // Same memory space and ownership constraints as 'a'
+    b: &r shrd gpu.global [i16; 16],
+    
+    // ERROR: This parameter should be declared as 'unq' (unique) instead of 'shrd' (shared)
+    // because it's used for assignment (c = a + b). Shared references are read-only and
+    // prevent data races by allowing multiple concurrent readers. Unique references are
+    // required for write operations to ensure exclusive access and prevent race conditions.
+    // The Descend compiler will detect this ownership violation and fail compilation.
+    c: &r shrd gpu.global [i16; 16]
+    
+// Execution context specification - defines how this function runs on GPU hardware
+// - grid: gpu.grid<X<1>, X<16>> - GPU execution grid with 1 block containing 16 threads
+// - The type system ensures GPU memory is only accessed in GPU execution contexts
+// - Prevents invalid cross-device memory accesses (CPU accessing GPU memory)
+) -[grid: gpu.grid<X<1>, X<16>>]-> () {
+    
+    // Vector addition operation - element-wise addition of arrays
+    // The compiler generates safe parallel code that:
+    // 1. Loads data from global memory to local memory for each thread
+    // 2. Performs vectorized addition using HIVM dialect operations
+    // 3. Stores results back to global memory safely
+    // The ownership system ensures this operation is race-free
+    //
+    // LAZY LOADING: Descend's compiler implements lazy loading strategies:
+    // - Memory loads are deferred until actually needed by computation
+    // - The HIVM dialect generates 'hivm.hir.load' operations that load from
+    //   global memory (gm) to local memory (ub) only when data is accessed
+    // - This minimizes memory bandwidth usage and improves cache efficiency
+    // - The type system ensures loads happen in the correct execution context
+    // - Shared references enable read-only access without unnecessary copies
+    c = a + b;
+    
+    // Unit return value - indicates successful completion
+    // In MLIR, this becomes a 'return' operation
+    ()
+}

src/codegen/mlir/builder/control_flow.rs

@@ -100,7 +100,7 @@ where
     let (then_region, true_value) = build_branch_region(case_true, ctx, build_expr)?;
 
     // Build the else region
-    let (else_region, false_value) = build_branch_region(case_false, ctx, build_expr)?;
+    let (else_region, _false_value) = build_branch_region(case_false, ctx, build_expr)?;
 
     // Determine result types based on whether branches produce values
     let result_types: Vec<Type> = if let Some(val) = true_value {

src/codegen/mlir/builder/place.rs

@@ -4,6 +4,7 @@ use super::super::error::MlirError;
 use super::context::MlirContext;
 use super::expr::build_expr;
 use crate::ast as desc;
+use crate::ast::{DataTyKind, Ownership};
 
 /// Build a place expression (variable lookup)
 pub fn build_place_expr<'ctx, 'a, 'b>(
@@ -57,6 +58,20 @@ where
 {
     use desc::PlaceExprKind;
 
+    // Check if we're assigning to a reference - if so, it must be unique
+    if let Some(ty) = &place_expr.ty {
+        if let desc::TyKind::Data(data_ty) = &ty.ty {
+            if let DataTyKind::Ref(ref_dty) = &data_ty.dty {
+                if ref_dty.own != Ownership::Uniq {
+                    return Err(MlirError::General(format!(
+                        "Assignment to non-unique reference is not allowed. Expected unique reference, found {:?}",
+                        ref_dty.own
+                    )));
+                }
+            }
+        }
+    }
+
     // Evaluate the right-hand side value
     let value = build_expr(value_expr, ctx)?
         .ok_or_else(|| MlirError::General("Missing value for assignment".to_string()))?;

src/codegen/mlir/mod.rs

@@ -72,22 +72,6 @@ fn build_module_internal(comp_unit: &CompilUnit) -> Result<String, MlirError> {
     Ok(builder.module().as_operation().to_string())
 }
 
-pub fn gen(comp_unit: &CompilUnit, _idx_checks: bool) -> String {
-    // Check if we need HIVM address spaces
-    if needs_hivm_address_space(comp_unit) {
-        to_mlir::types::generate_mlir_string_with_hivm(comp_unit)
-    } else {
-        // Use internal helper, but handle errors by falling back to string generation
-        match build_module_internal(comp_unit) {
-            Ok(mlir_string) => mlir_string,
-            Err(_) => {
-                // Fallback to string generation if internal building fails
-                to_mlir::types::generate_mlir_string_with_hivm(comp_unit)
-            }
-        }
-    }
-}
-
 pub fn gen_checked(comp_unit: &CompilUnit, _idx_checks: bool) -> Result<String, String> {
     // Check if we need HIVM address spaces
     if needs_hivm_address_space(comp_unit) {

src/codegen/mlir/to_mlir/types.rs

@@ -1,4 +1,4 @@
-use crate::ast::{AtomicTy, BaseExec, DataTy, DataTyKind, FunDef, Memory, Nat, NatCtx, ScalarTy, Ty, TyKind};
+use crate::ast::{AtomicTy, BaseExec, DataTy, DataTyKind, FunDef, Memory, Nat, NatCtx, Ownership, ScalarTy, Ty, TyKind};
 use melior::{
     dialect::func,
     ir::{
@@ -23,35 +23,6 @@ fn nat_to_dimension(nat: &Nat) -> String {
     }
 }
 
-/// Helper function to create HACC attributes for GPU functions
-/// This function creates proper MLIR attribute objects for HACC attributes.
-/// Note: This requires the HACC dialect to be registered in the MLIR context.
-fn create_hacc_attributes<'c>(context: &'c Context) -> Vec<(Identifier<'c>, Attribute<'c>)> {
-    // Create HACC entry attribute (hacc.entry) - this is a unit attribute
-    let entry_attr = Attribute::parse(context, "unit")
-        .expect("Failed to create HACC entry attribute");
-    
-    // Create HACC function type attribute (hacc.function_kind = DEVICE)
-    // Note: This will fail if the HACC dialect is not registered in the context
-    let func_type_attr = Attribute::parse(context, "#hacc.function_kind<DEVICE>")
-        .expect("Failed to create HACC function type attribute - ensure HACC dialect is registered");
-    
-    vec![
-        (Identifier::new(context, "hacc.entry"), entry_attr),
-        (Identifier::new(context, "hacc.function_kind"), func_type_attr),
-    ]
-}
-
-/// Helper function to generate HACC attributes string for GPU functions
-/// This function generates the MLIR string representation of HACC attributes
-/// for use in string-based MLIR generation. The proper attribute objects
-/// are available through create_hacc_attributes() when the HACC dialect is registered.
-fn generate_hacc_attributes_string(context: &Context) -> String {
-    // Generate the string representation directly since we know the exact format
-    // The create_hacc_attributes() function is available for when the HACC dialect is registered
-    " attributes {hacc.entry, hacc.function_kind = #hacc.function_kind<DEVICE>}".to_string()
-}
-
 /// Helper function to convert ScalarTy to MLIR Type
 fn scalar_ty_to_mlir<'c>(scalar_ty: &ScalarTy, context: &'c Context) -> Type<'c> {
     match scalar_ty {
@@ -505,73 +476,6 @@ fn collect_parameter_usage(fun: &crate::ast::FunDef) -> std::collections::HashMa
     param_usage
 }
 
-/// Collect which parameters are referenced in the function body (legacy function for compatibility)
-fn collect_used_parameters(fun: &crate::ast::FunDef) -> std::collections::HashSet<String> {
-    use crate::ast::{Expr, ExprKind, PlaceExprKind};
-    use std::collections::HashSet;
-    
-    let mut used_params = HashSet::new();
-    
-    fn walk_expr(expr: &Expr, used_params: &mut HashSet<String>) {
-        match &expr.expr {
-            ExprKind::PlaceExpr(place_expr) => {
-                if let PlaceExprKind::Ident(ident) = &place_expr.pl_expr {
-                    used_params.insert(ident.name.to_string());
-                }
-            }
-            ExprKind::BinOp(_, lhs, rhs) => {
-                walk_expr(lhs, used_params);
-                walk_expr(rhs, used_params);
-            }
-            ExprKind::Let(_, _, value_expr) => {
-                walk_expr(value_expr, used_params);
-            }
-            ExprKind::Seq(exprs) => {
-                for expr in exprs {
-                    walk_expr(expr, used_params);
-                }
-            }
-            ExprKind::Assign(place_expr, value_expr) => {
-                if let PlaceExprKind::Ident(ident) = &place_expr.pl_expr {
-                    used_params.insert(ident.name.to_string());
-                }
-                walk_expr(value_expr, used_params);
-            }
-            ExprKind::App(_, _, args) => {
-                for arg in args {
-                    walk_expr(arg, used_params);
-                }
-            }
-            ExprKind::IfElse(cond, case_true, case_false) => {
-                walk_expr(cond, used_params);
-                walk_expr(case_true, used_params);
-                walk_expr(case_false, used_params);
-            }
-            ExprKind::If(cond, case_true) => {
-                walk_expr(cond, used_params);
-                walk_expr(case_true, used_params);
-            }
-            ExprKind::ForNat(_, _, body) => {
-                walk_expr(body, used_params);
-            }
-            ExprKind::Ref(_, _, place_expr) => {
-                if let PlaceExprKind::Ident(ident) = &place_expr.pl_expr {
-                    used_params.insert(ident.name.to_string());
-                }
-            }
-            ExprKind::Unsafe(expr) => {
-                walk_expr(expr, used_params);
-            }
-            _ => {
-                // Other expression types don't contain variable references
-            }
-        }
-    }
-    
-    walk_expr(&fun.body.body, &mut used_params);
-    used_params
-}
-
 /// Generate body operations for GPU functions
 fn generate_body_operations(
     fun: &crate::ast::FunDef,
@@ -680,6 +584,18 @@ fn generate_body_operations(
                     // Find the parameter declaration to get its type and index
                     if let Some((param_idx, param_decl)) = fun.param_decls.iter().enumerate().find(|(_, p)| p.ident.name == ident.name) {
                         if let Some(param_ty) = &param_decl.ty {
+                            // Check if we're assigning to a reference - if so, it must be unique
+                            if let TyKind::Data(data_ty) = &param_ty.ty {
+                                if let DataTyKind::Ref(ref_dty) = &data_ty.dty {
+                                    if ref_dty.own != Ownership::Uniq {
+                                        panic!(
+                                            "Assignment to non-unique reference is not allowed. Expected unique reference, found {:?}",
+                                            ref_dty.own
+                                        );
+                                    }
+                                }
+                            }
+                            
                             // Generate the target parameter type (should be gm address space)
                             let target_type = get_mlir_type_string_with_address_space(param_ty, context);
 
@@ -856,7 +772,7 @@ fn generate_function_with_body(fun: &crate::ast::FunDef, context: &Context) -> S
         // TODO: When HACC dialect is registered in the MLIR context, replace this with:
         // let hacc_attributes = create_hacc_attributes(context);
         // and use the attributes with MLIR operation builders instead of string generation
-        signature.push_str(&generate_hacc_attributes_string(context));
+        signature.push_str(" attributes {hacc.entry, hacc.function_kind = #hacc.function_kind<DEVICE>}");
     }
 
     signature.push_str(" {\n");

src/lib.rs

@@ -29,18 +29,4 @@ pub fn compile(file_path: &str, backend: Backend) -> Result<(String, String), Er
     let ast_string = format!("{:#?}", compil_unit.items);
 
     Ok((code_string, ast_string))
-}
-
-pub fn compile_unchecked(file_path: &str, backend: Backend) -> Result<String, ErrorReported> {
-    let source = parser::SourceCode::from_file(file_path)?;
-    let mut compil_unit = parser::parse(&source)?;
-
-    ty_check::ty_check(&mut compil_unit)?;
-
-    let code_string = match backend {
-        Backend::Cuda => codegen::cuda::gen(&compil_unit, false),
-        Backend::Mlir => codegen::mlir::gen(&compil_unit, false),
-    };
-
-    Ok(code_string)
-}
+}

src/main.rs

@@ -21,10 +21,6 @@ struct Args {
     /// Print Ast
     #[arg(short, long)]
     print_ast: bool,
-
-    /// Skip MLIR verification checks
-    #[arg(long)]
-    no_checks: bool,
 }
 
 /// Backend selection passed via CLI
@@ -52,23 +48,13 @@ fn main() {
     let output_dir = &args.output_dir;
 
     // Compile using Descend
-    let (code_string, ast_string) = if args.no_checks {
-        let code_string = match descend::compile_unchecked(&input_path.to_string_lossy(), backend) {
+    let (code_string, ast_string) = match descend::compile(&input_path.to_string_lossy(), backend) {
             Ok(output) => output,
             Err(e) => {
                 eprintln!("Compilation failed: {:?}", e);
                 std::process::exit(1);
             }
-        };
-        (code_string, String::new())
-    } else {
-        match descend::compile(&input_path.to_string_lossy(), backend) {
-            Ok(output) => output,
-            Err(e) => {
-                eprintln!("Compilation failed: {:?}", e);
-                std::process::exit(1);
-            }
-        }
+        
     };
 
     // Generate output file path with appropriate extension based on backend

tests/mlir.rs

@@ -1,4 +1,7 @@
 #[path = "mlir/core.rs"]
 mod core;
 
+#[path = "mlir/error_examples.rs"]
+mod error_examples;
+
 const BACKEND: descend::Backend = descend::Backend::Mlir;

tests/mlir/core.rs

@@ -1,4 +1,2 @@
-type Res = Result<(), descend::error::ErrorReported>;
-
 // Automatically generate tests for all .desc files in examples/core/
 descend_derive::generate_desc_tests!();