Fix dead code elimination for Expr::Slice

rose-lang · Nov 16, 2023 · ff9ea61 · ff9ea61
1 parent 16fe5e6
commit ff9ea61
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/crates/wasm/src/lib.rs b/crates/wasm/src/lib.rs
@@ -139,7 +139,10 @@ impl<'a, O: Eq + Hash, T: Refs<'a, Opaque = O>> Topsort<'a, O, T> {
                     vars.follow(index, instr.var);
                 }
                 &Expr::Member { tuple, .. } => vars.follow(tuple, instr.var),
-                &Expr::Slice { .. } => vars.live(instr.var),
+                &Expr::Slice { index, .. } => {
+                    vars.live(instr.var);
+                    vars.follow(index, instr.var);
+                }
                 &Expr::Field { .. } => vars.live(instr.var),
                 &Expr::Unary { arg, .. } => vars.follow(arg, instr.var),
                 &Expr::Binary { left, right, .. } => {

diff --git a/packages/core/src/index.test.ts b/packages/core/src/index.test.ts
@@ -934,4 +934,12 @@ describe("valid", () => {
     const h = await compile(g);
     expect(h()).toEqual([[0]]);
   });
+
+  test("compile gradient with dynamic index", async () => {
+    const T = struct({ v: Vec(1, Real), i: 1 });
+    const f = fn([T], Real, ({ v, i }) => v[i]);
+    const g = fn([T], T, (x) => vjp(f)(x).grad(1));
+    const h = await compile(g);
+    expect(h({ v: [2], i: 0 })).toEqual({ v: [1], i: 0 });
+  });
 });