zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-25 01:41:17 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix bug in tape order during simplify (#816)

Browse Source 
* fix bug in tape order during simplify

* properly fix compile

* last bug
This commit is contained in:
Awni Hannun 2024-03-11 17:29:05 -07:00 committed by GitHub
parent 0ae22b915b
commit 0e95b64942
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 55 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
mlx/compile.cpp
View File

@ -415,6 +415,11 @@ void compile_simplify(
}
tape = std::move(new_tape);
std::unordered_map<std::uintptr_t, uint32_t> tape_order;
for (uint32_t i = 0; i < tape.size(); ++i) {
tape_order.insert({tape[i].id(), i});
}
std::unordered_set<uintptr_t> output_set;
for (auto& o : outputs) {
output_set.insert(o.id());
@ -437,17 +442,23 @@ void compile_simplify(
if (mask[j]) {
continue;
}
auto& src = parents->second[j].first;
auto& dst = parents->second[i].first;
auto src_idx = j;
auto dst_idx = i;
if (tape_order[parents->second[src_idx].first.id()] <
tape_order[parents->second[dst_idx].first.id()]) {
std::swap(src_idx, dst_idx);
}
auto& src = parents->second[src_idx].first;
auto& dst = parents->second[dst_idx].first;
if (src.id() != dst.id() && array_equivalent(src, dst) &&
output_set.find(src.id()) == output_set.end()) {
merge(dst, src, parents_map);
mask[j] = true;
mask[src_idx] = true;
}
}
}
// Erase orphaned parents so we don't keep fusing with them
for (int i = N - 1; i > 0; --i) {
for (int i = N - 1; i >= 0; --i) {
if (mask[i]) {
parents->second.erase(parents->second.begin() + i);
}

40
python/tests/test_compile.py
View File

@ -613,6 +613,46 @@ class TestCompile(mlx_tests.MLXTestCase):
out = fun()
mx.eval(out)
def test_compile_vjp(self):
def fun(w):
w1 = w + w
w2 = w + w
return w @ w1 + w2 @ w2
def step(w):
out, grad = mx.vjp(fun, (w,), (mx.array([[1.0, 1.0], [1.0, 1.0]]),))
return out[0], grad[0]
w = mx.zeros((2, 2))
mx.eval(w)
expected = step(w)
out = mx.compile(step)(w)
self.assertTrue(mx.allclose(expected[0], out[0]))
self.assertTrue(mx.allclose(expected[1], out[1]))
def fun(w1, w2, x):
x = x @ w1
y = x @ w2
x = x + y * y
return (x * x).sum()
w1 = mx.zeros((4, 4))
w2 = mx.zeros((4, 4))
x = mx.zeros((4, 4))
def step(w1, w2, x):
loss, gradient = mx.value_and_grad(fun)(w1, w2, x)
w1 = w1 + gradient
return loss, w1
mx.eval(x, w1, w2)
expected = step(w1, w2, x)
out = mx.compile(step)(w1, w2, x)
self.assertTrue(mx.allclose(expected[0], out[0]))
self.assertTrue(mx.allclose(expected[1], out[1]))
if __name__ == "__main__":
unittest.main()