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

recompile when shapeless is different (#1776)

Browse Source
This commit is contained in:
Awni Hannun
2025-01-20 21:07:10 -08:00
committed by GitHub
parent a8666a757a
commit 90532b1f37
2 changed files with 11 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
mlx/compile.cpp
View File

@@ -197,8 +197,10 @@ std::uintptr_t get_function_address(const std::function<T(U...)>& fun) {
class CompilerCache {
public:
struct CacheEntry {
CacheEntry(Stream stream) : stream(stream) {};
CacheEntry(Stream stream, bool shapeless)
: stream(stream), shapeless(shapeless) {};
Stream stream;
bool shapeless;
std::vector<array> inputs;
std::vector<array> outputs;
std::vector<array> tape;
@@ -245,6 +247,9 @@ class CompilerCache {
if (entry.stream != stream) {
continue;
}
if (entry.shapeless != shapeless) {
continue;
}
// Check the inputs match and return if so
if (has_same_shape_and_dtype(inputs, entry.inputs) &&
@@ -253,7 +258,7 @@ class CompilerCache {
}
}
// Otherwise append a new cache entry
entries.push_back(CacheEntry{stream});
entries.push_back(CacheEntry{stream, shapeless});
return entries.back();
}