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

throw for certain cases of non captured inputs in compile (#1401)

Browse Source
This commit is contained in:
Awni Hannun
2024-09-09 14:54:31 -07:00
committed by GitHub
parent dc627dcb5e
commit 3ae6aabe9f
5 changed files with 70 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
mlx/compile.cpp
View File

@@ -306,21 +306,27 @@ std::pair<std::vector<array>, std::vector<array>> compile_trace(
// Traverses the graph to build a tape and a map of array ids to their parents
std::pair<std::vector<array>, ParentsMap> compile_dfs(
const std::vector<array>& inputs,
const std::vector<array>& outputs) {
const std::vector<array>& outputs,
const std::vector<array>& original_inputs) {
std::function<void(const array&)> recurse;
std::vector<array> tape;
std::unordered_set<std::uintptr_t> input_set;
std::unordered_set<std::uintptr_t> original_input_set;
std::unordered_map<std::uintptr_t, std::vector<std::pair<array, int>>>
parents_map;
for (int i = 0; i < inputs.size(); ++i) {
auto in = inputs[i];
input_set.insert(in.id());
input_set.insert(inputs[i].id());
original_input_set.insert(original_inputs[i].id());
}
// DFS the graph to build the tape, and log parents and scalars
std::unordered_set<std::uintptr_t> cache;
recurse = [&](const array& a) {
auto id = a.id();
if (original_input_set.find(id) != original_input_set.end()) {
throw std::invalid_argument(
"[compile] Attempting to compile a function with uncaptured inputs is not allowed.");
}
if (cache.find(id) != cache.end()) {
return;
}
@@ -833,7 +839,7 @@ std::function<std::vector<array>(const std::vector<array>&)> compile(
std::unordered_map<uintptr_t, std::vector<std::pair<array, int>>>
parents_map;
std::tie(entry.tape, parents_map) =
compile_dfs(entry.inputs, entry.outputs);
compile_dfs(entry.inputs, entry.outputs, inputs);
// Simplify the tape
if (compile_mode() != CompileMode::no_simplify) {

2
mlx/fast.cpp
View File

@@ -972,7 +972,7 @@ void write_signature(
{"threadgroups_per_grid", "uint3"},
{"threads_per_grid", "uint3"},
{"threads_per_simdgroup", "uint"},
{"thread_per_threadgroup", "uint3"},
{"threads_per_threadgroup", "uint3"},
};
std::vector<std::pair<std::string, std::string>> attrs;
for (const auto& [attr, dtype] : metal_attributes) {