dnrops
/
burn
mirror of https://github.com/tracel-ai/burn.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Force constant vec indexing

This commit is contained in:
nathaniel 2024-04-22 16:00:00 -04:00
parent 6c708527b9
commit 85264532d1
1 changed files with 74 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
crates/burn-wgpu/src/compiler/wgsl/instructions.rs
View File

@ -249,6 +249,18 @@ impl Display for Instruction {
}
Instruction::Not { input, out } => f.write_fmt(format_args!("{out} = !{input};\n")),
Instruction::Index { lhs, rhs, out } => {
if let Variable::Local {
index: _,
item: _,
scope_depth: _,
} = out
{
match rhs {
Variable::GlobalScalar(_, _, _) => todo!(),
_ => panic!("Only constant indexing is supported, got {:?}", rhs),
}
};
let item = out.item();
f.write_fmt(format_args!("{out} = {item}({lhs}[{rhs}]);\n"))
}
@ -379,7 +391,20 @@ for (var {i}: u32 = {start}; {i} < {end}; {i}++) {{
f.write_str("}\n")
}
Instruction::IndexAssign { lhs, rhs, out } => match lhs.item() {
Instruction::IndexAssign { lhs, rhs, out } => {
if let Variable::Local {
index: _,
item: _,
scope_depth: _,
} = out
{
match lhs {
Variable::GlobalScalar(_, _, _) => todo!(),
_ => panic!("Only constant indexing is supported, got {:?}", lhs),
}
};
match lhs.item() {
Item::Vec4(elem) => {
let lhs0 = lhs.index(0);
let lhs1 = lhs.index(1);
@ -429,7 +454,8 @@ for (var {i}: u32 = {start}; {i} < {end}; {i}++) {{
};
f.write_fmt(format_args!("{out}[{lhs}] = {casting_type}({rhs});\n"))
}
},
}
}
Instruction::If { cond, instructions } => {
f.write_fmt(format_args!("if {cond} {{\n"))?;
for i in instructions {