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

Kernel generation (#614)

Browse Source 
Generate reusable element-wise kernels given a computation graph.
This commit is contained in:
Angelos Katharopoulos
2024-02-07 13:15:59 -08:00
committed by GitHub
parent 5fd11c347d
commit 28eac18571
19 changed files with 1302 additions and 459 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
mlx/graph_utils.cpp
View File

@@ -12,27 +12,24 @@
namespace mlx::core {
struct NodeNamer {
std::unordered_map<std::uintptr_t, std::string> names;
std::string get_name(const array& x) {
auto it = names.find(x.id());
if (it == names.end()) {
// Get the next name in the sequence
// [A, B, ..., Z, AA, AB, ...]
std::vector<char> letters;
auto var_num = names.size() + 1;
while (var_num > 0) {
letters.push_back('A' + (var_num - 1) % 26);
var_num = (var_num - 1) / 26;
}
std::string name(letters.rbegin(), letters.rend());
names.insert({x.id(), name});
return name;
const std::string& NodeNamer::get_name(const array& x) {
auto it = names.find(x.id());
if (it == names.end()) {
// Get the next name in the sequence
// [A, B, ..., Z, AA, AB, ...]
std::vector<char> letters;
auto var_num = names.size() + 1;
while (var_num > 0) {
letters.push_back('A' + (var_num - 1) % 26);
var_num = (var_num - 1) / 26;
}
return it->second;
std::string name(letters.rbegin(), letters.rend());
names.insert({x.id(), name});
return get_name(x);
}
};
return it->second;
}
void depth_first_traversal(
std::function<void(array)> callback,