feat: handle dynamic arrays

Author: lucarlig

examples/core/vec_add_threads_sched.desc

@@ -0,0 +1,23 @@
+fn gpu_vec_add<r: prv, n: nat>(
+    a_array: &r uniq gpu.global [i32; n],
+    b_array: &r shrd gpu.global [i32; n]
+) -[grid: gpu.grid<X<64>, X<1024>>]-> () {
+    sched block in grid {
+        sched thread in block {
+            let a_ref = &uniq (*a_array).to_view.grp::<1024>[[block]][[thread]];
+            *a_ref = *a_ref + (*b_array).to_view.grp::<1024>[[block]][[thread]]
+        }
+    }
+}
+
+fn main() -[t: cpu.thread]-> () <'a>{
+    let a_array: &'a uniq cpu.mem [i32; 64*1024] = unsafe _;
+    let b_array: &'a shrd cpu.mem [i32; 64*1024] = unsafe _;
+    let mut gpu = gpu_device(0);
+
+    let mut a_gpu_array = gpu_alloc_copy(&uniq gpu, &shrd *a_array);
+    let b_gpu_array = gpu_alloc_copy(&uniq gpu, &shrd *b_array);
+    gpu_vec_add::<<<X<64>, X<1024>>>>(&uniq a_gpu_array, &shrd b_gpu_array);
+    copy_to_host(&shrd a_gpu_array, a_array);
+    ()
+}

src/codegen/mlir/to_mlir/types.rs

@@ -17,14 +17,12 @@ pub trait ToMlir {
 }
 
 impl Nat {
-    fn to_dimension_i64(self: &Self) -> i64 {
+    fn to_dimension(self: &Self) -> Vec<i64> {
         // Try to evaluate the Nat with an empty context
         let nat_ctx = NatCtx::new();
         match self.eval(&nat_ctx) {
-            Ok(size) => size as i64,
-            Err(_) => panic!(
-                "Array dimensions must be compile-time known. Dynamic arrays are not supported."
-            ),
+            Ok(size) => vec![size as i64],
+            Err(_) => vec![],
         }
     }
 }
@@ -149,8 +147,8 @@ fn ref_array_to_mlir<'c>(
 ) -> Type<'c> {
     // Array reference -> memref with dimensions
     let elem_type = elem_ty.to_mlir(context);
-    let dim = size.to_dimension_i64();
-    let base_type: Type<'c> = MemRefType::new(elem_type, &[dim], None, None).into();
+    let dim = size.to_dimension();
+    let base_type: Type<'c> = MemRefType::new(elem_type, &dim, None, None).into();
 
     // Add HIVM address space if needed
     let base_str = base_type.to_string();
@@ -169,8 +167,8 @@ fn ref_at_to_mlir<'c>(inner: &DataTy, mem: &Memory, context: &'c Context) -> Typ
         }
         DataTyKind::Array(elem_ty, size) | DataTyKind::ArrayShape(elem_ty, size) => {
             let elem_type = elem_ty.to_mlir(context);
-            let dim = size.to_dimension_i64();
-            MemRefType::new(elem_type, &[dim], None, None).into()
+            let dim = size.to_dimension();
+            MemRefType::new(elem_type, &dim, None, None).into()
         }
         DataTyKind::Tuple(_) => {
             unimplemented!("Tuple references with At not yet supported in MLIR conversion")
@@ -239,13 +237,13 @@ impl ToMlir for DataTy {
             DataTyKind::Ident(ident) => ident.to_mlir(context),
             DataTyKind::Array(elem_ty, size) => {
                 let elem_type = elem_ty.to_mlir(context);
-                let dim = size.to_dimension_i64();
-                MemRefType::new(elem_type, &[dim], None, None).into()
+                let dim = size.to_dimension();
+                MemRefType::new(elem_type, &dim, None, None).into()
             }
             DataTyKind::ArrayShape(elem_ty, size) => {
                 let elem_type = elem_ty.to_mlir(context);
-                let dim = size.to_dimension_i64();
-                MemRefType::new(elem_type, &[dim], None, None).into()
+                let dim = size.to_dimension();
+                MemRefType::new(elem_type, &dim, None, None).into()
             }
             DataTyKind::Struct(_) => {
                 unimplemented!("Struct types not yet supported in MLIR conversion")

tests/mlir/snapshots/mlir__core__vec_add_chain.snap

@@ -0,0 +1,20 @@
+---
+source: tests/mlir/core.rs
+expression: output
+---
+module {
+  func.func @add(%arg0: memref<16xi16, #hivm.address_space<gm>>, %arg1: memref<16xi16, #hivm.address_space<gm>>, %arg2: memref<16xi16, #hivm.address_space<gm>>, %arg3: memref<16xi16, #hivm.address_space<gm>>) attributes {hacc.entry, hacc.function_kind = #hacc.function_kind<DEVICE>} {
+    %alloc = memref.alloc() : memref<16xi16, #hivm.address_space<ub>>
+    hivm.hir.load ins(%arg0 : memref<16xi16, #hivm.address_space<gm>>) outs(%alloc : memref<16xi16, #hivm.address_space<ub>>)
+    %alloc_0 = memref.alloc() : memref<16xi16, #hivm.address_space<ub>>
+    hivm.hir.load ins(%arg1 : memref<16xi16, #hivm.address_space<gm>>) outs(%alloc_0 : memref<16xi16, #hivm.address_space<ub>>)
+    %alloc_1 = memref.alloc() : memref<16xi16, #hivm.address_space<ub>>
+    hivm.hir.load ins(%arg2 : memref<16xi16, #hivm.address_space<gm>>) outs(%alloc_1 : memref<16xi16, #hivm.address_space<ub>>)
+    %alloc_2 = memref.alloc() : memref<16xi16, #hivm.address_space<ub>>
+    hivm.hir.vadd ins(%alloc, %alloc_0 : memref<16xi16, #hivm.address_space<ub>>, memref<16xi16, #hivm.address_space<ub>>) outs(%alloc_2 : memref<16xi16, #hivm.address_space<ub>>)
+    %alloc_3 = memref.alloc() : memref<16xi16, #hivm.address_space<ub>>
+    hivm.hir.vadd ins(%alloc_2, %alloc_1 : memref<16xi16, #hivm.address_space<ub>>, memref<16xi16, #hivm.address_space<ub>>) outs(%alloc_3 : memref<16xi16, #hivm.address_space<ub>>)
+    hivm.hir.store ins(%alloc_3 : memref<16xi16, #hivm.address_space<ub>>) outs(%arg3 : memref<16xi16, #hivm.address_space<gm>>)
+    return
+  }
+}