From 39e2ea9372530618a5f6aab8cf6c8b371849b637 Mon Sep 17 00:00:00 2001
From: Aaron Councilman <aaronjc4@illinois.edu>
Date: Fri, 28 Feb 2025 11:16:08 -0600
Subject: [PATCH] Allow let-tuple destruction in scheduling language

---
 juno_samples/matmul/src/cpu.sch |  5 ++---
 juno_scheduler/src/compile.rs   | 18 ++++++++++++++++++
 juno_scheduler/src/lang.y       |  3 +++
 3 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/juno_samples/matmul/src/cpu.sch b/juno_samples/matmul/src/cpu.sch
index 7577aebb..69f1811d 100644
--- a/juno_samples/matmul/src/cpu.sch
+++ b/juno_samples/matmul/src/cpu.sch
@@ -46,16 +46,15 @@ associative(matmul@outer);
 // Parallelize by computing output array as 16 chunks
 let par = matmul@outer \ matmul@inner;
 fork-tile![4](par);
-let res = fork-reshape[[1, 3], [0], [2]](par); let outer = res.0; let inner = res.1;
+let (outer, inner, _) = fork-reshape[[1, 3], [0], [2]](par);
 parallelize!(outer \ inner);
 
 let body = outline(inner);
 cpu(body);
 
 // Tile for cache, assuming 64B cache lines
-fork-split(body);
 fork-tile![16](body);
-let res = fork-reshape[[0, 2, 4, 1, 3], [5]](body); let outer = res.0; let inner = res.1;
+let (outer, inner) = fork-reshape[[0, 2, 4, 1, 3], [5]](body);
 
 reduce-slf(inner);
 unforkify!(body);
diff --git a/juno_scheduler/src/compile.rs b/juno_scheduler/src/compile.rs
index 0f473f0d..51ba3699 100644
--- a/juno_scheduler/src/compile.rs
+++ b/juno_scheduler/src/compile.rs
@@ -207,6 +207,24 @@ fn compile_stmt(
                 exp: compile_exp_as_expr(expr, lexer, macrostab, macros)?,
             }])
         }
+        parser::Stmt::LetsStmt { span: _, vars, expr } => {
+            let tmp = format!("{}_tmp", macros.uniq());
+            Ok(std::iter::once(ir::ScheduleStmt::Let {
+                var: tmp.clone(),
+                exp: compile_exp_as_expr(expr, lexer, macrostab, macros)?,
+            }).chain(vars.into_iter().enumerate()
+                .map(|(idx, v)| {
+                    let var = lexer.span_str(v).to_string();
+                    ir::ScheduleStmt::Let {
+                        var,
+                        exp: ir::ScheduleExp::TupleField {
+                            lhs: Box::new(ir::ScheduleExp::Variable { var: tmp.clone() }),
+                            field: idx,
+                        }
+                    }
+                })
+            ).collect())
+        }
         parser::Stmt::AssignStmt { span: _, var, rhs } => {
             let var = lexer.span_str(var).to_string();
             Ok(vec![ir::ScheduleStmt::Assign {
diff --git a/juno_scheduler/src/lang.y b/juno_scheduler/src/lang.y
index 69d50014..451f035b 100644
--- a/juno_scheduler/src/lang.y
+++ b/juno_scheduler/src/lang.y
@@ -19,6 +19,8 @@ Schedule -> OperationList
 Stmt -> Stmt
   : 'let' 'ID' '=' Expr ';'
       { Stmt::LetStmt { span: $span, var: span_of_tok($2), expr: $4 } }
+  | 'let' '(' Ids ')' '=' Expr ';'
+      { Stmt::LetsStmt { span: $span, vars: rev($3), expr: $6 } }
   | 'ID' '=' Expr ';'
       { Stmt::AssignStmt { span: $span, var: span_of_tok($1), rhs: $3 } }
   | Expr ';'
@@ -157,6 +159,7 @@ pub enum OperationList {
 
 pub enum Stmt {
   LetStmt    { span: Span, var: Span, expr: Expr },
+  LetsStmt   { span: Span, vars: Vec<Span>, expr: Expr },
   AssignStmt { span: Span, var: Span, rhs: Expr },
   ExprStmt   { span: Span, exp: Expr },
   Fixpoint   { span: Span, limit: FixpointLimit, body: Box<OperationList> },
-- 
GitLab