zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 09:21:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

compile changes if stream changes (#1644)

Browse Source
This commit is contained in:
Awni Hannun 2024-12-03 14:37:44 -08:00 committed by GitHub
parent 9d40e521d7
commit e047fd977d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 24 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
mlx/compile.cpp
View File

@ -211,6 +211,8 @@ std::uintptr_t get_function_address(const std::function<T(U...)>& fun) {
class CompilerCache {
public:
struct CacheEntry {
CacheEntry(Stream stream) : stream(stream) {};
Stream stream;
std::vector<array> inputs;
std::vector<array> outputs;
std::vector<array> tape;
@ -227,6 +229,7 @@ class CompilerCache {
const std::vector<uint64_t>& constants) {
// Find the cache entries for |fun_id|.
std::vector<CacheEntry>& entries = cache_[fun_id];
// Compare if 2 arrays have same shape and dtype.
auto has_same_shape_and_dtype = [shapeless](
const std::vector<array>& in1,
@ -247,11 +250,16 @@ class CompilerCache {
}
return true;
};
// Loop over entries and check inputs match i.e. shapes and types must be
// equal. Note this could get really slow if one compiles the same
// function with many different shapes. May want to store entries in a
// more easily searchable structure.
// Loop over entries and check:
// - Default stream and device match the entry's default stream
// - Inputs match i.e. shapes and types must be equal.
auto stream = default_stream(default_device());
for (CacheEntry& entry : entries) {
// Check that the default stream and device match
if (entry.stream != stream) {
continue;
}
// Check the inputs match and return if so
if (has_same_shape_and_dtype(inputs, entry.inputs) &&
constants == entry.constants) {
@ -259,7 +267,7 @@ class CompilerCache {
}
}
// Otherwise append a new cache entry
entries.push_back(CacheEntry{});
entries.push_back(CacheEntry{stream});
return entries.back();
}

8
python/src/array.cpp
View File

@ -184,10 +184,6 @@ void init_array(nb::module_& m) {
R"pbdoc(
A helper object to apply updates at specific indices.
)pbdoc")
.def(
nb::init<const array&>(),
"x"_a,
nb::sig("def __init__(self, x: array)"))
.def("__getitem__", &ArrayAt::set_indices, "indices"_a.none())
.def("add", &ArrayAt::add, "value"_a)
.def("subtract", &ArrayAt::subtract, "value"_a)
@ -202,10 +198,6 @@ void init_array(nb::module_& m) {
R"pbdoc(
A helper object to iterate over the 1st dimension of an array.
)pbdoc")
.def(
nb::init<const array&>(),
"x"_a,
nb::sig("def __init__(self, x: array)"))
.def("__next__", &ArrayPythonIterator::next)
.def("__iter__", [](const ArrayPythonIterator& it) { return it; });

1
python/src/stream.cpp
View File

@ -48,7 +48,6 @@ void init_stream(nb::module_& m) {
R"pbdoc(
A stream for running operations on a given device.
)pbdoc")
.def(nb::init<int, Device>(), "index"_a, "device"_a)
.def_ro("device", &Stream::device)
.def(
"__repr__",

11
tests/compile_tests.cpp
View File

@ -719,3 +719,14 @@ TEST_CASE("test compile strides") {
CHECK_EQ(out.strides().size(), 3);
}
}
TEST_CASE("test compile change streams") {
auto cfun = compile(simple_fun);
auto out = cfun({array(1.0f), array(2.0f)})[0];
CHECK_EQ(out.primitive().stream(), default_stream(default_device()));
auto s = new_stream(default_device());
StreamContext sctx(s);
out = cfun({array(1.0f), array(2.0f)})[0];
CHECK_EQ(out.primitive().stream(), s);
}