Skip to content
Snippets Groups Projects
Commit 39e2ea93 authored by Aaron Councilman's avatar Aaron Councilman
Browse files

Allow let-tuple destruction in scheduling language

parent 2d577f79
No related branches found
No related tags found
1 merge request!205Fork reshape
Pipeline #201929 passed
...@@ -46,16 +46,15 @@ associative(matmul@outer); ...@@ -46,16 +46,15 @@ associative(matmul@outer);
// Parallelize by computing output array as 16 chunks // Parallelize by computing output array as 16 chunks
let par = matmul@outer \ matmul@inner; let par = matmul@outer \ matmul@inner;
fork-tile![4](par); 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); parallelize!(outer \ inner);
let body = outline(inner); let body = outline(inner);
cpu(body); cpu(body);
// Tile for cache, assuming 64B cache lines // Tile for cache, assuming 64B cache lines
fork-split(body);
fork-tile![16](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); reduce-slf(inner);
unforkify!(body); unforkify!(body);
......
...@@ -207,6 +207,24 @@ fn compile_stmt( ...@@ -207,6 +207,24 @@ fn compile_stmt(
exp: compile_exp_as_expr(expr, lexer, macrostab, macros)?, 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 } => { parser::Stmt::AssignStmt { span: _, var, rhs } => {
let var = lexer.span_str(var).to_string(); let var = lexer.span_str(var).to_string();
Ok(vec![ir::ScheduleStmt::Assign { Ok(vec![ir::ScheduleStmt::Assign {
......
...@@ -19,6 +19,8 @@ Schedule -> OperationList ...@@ -19,6 +19,8 @@ Schedule -> OperationList
Stmt -> Stmt Stmt -> Stmt
: 'let' 'ID' '=' Expr ';' : 'let' 'ID' '=' Expr ';'
{ Stmt::LetStmt { span: $span, var: span_of_tok($2), expr: $4 } } { 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 ';' | 'ID' '=' Expr ';'
{ Stmt::AssignStmt { span: $span, var: span_of_tok($1), rhs: $3 } } { Stmt::AssignStmt { span: $span, var: span_of_tok($1), rhs: $3 } }
| Expr ';' | Expr ';'
...@@ -157,6 +159,7 @@ pub enum OperationList { ...@@ -157,6 +159,7 @@ pub enum OperationList {
pub enum Stmt { pub enum Stmt {
LetStmt { span: Span, var: Span, expr: Expr }, LetStmt { span: Span, var: Span, expr: Expr },
LetsStmt { span: Span, vars: Vec<Span>, expr: Expr },
AssignStmt { span: Span, var: Span, rhs: Expr }, AssignStmt { span: Span, var: Span, rhs: Expr },
ExprStmt { span: Span, exp: Expr }, ExprStmt { span: Span, exp: Expr },
Fixpoint { span: Span, limit: FixpointLimit, body: Box<OperationList> }, Fixpoint { span: Span, limit: FixpointLimit, body: Box<OperationList> },
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment