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

Add types and order in kernel name (#831)

Browse Source
This commit is contained in:
Angelos Katharopoulos 2024-03-13 20:34:06 -07:00 committed by GitHub
parent 43abc402d8
commit 1efee9db09
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 15 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
mlx/backend/common/compiled.cpp
View File

@ -1,6 +1,7 @@
// Copyright © 2023-2024 Apple Inc. // Copyright © 2023-2024 Apple Inc.
#include "mlx/backend/common/compiled.h" #include "mlx/backend/common/compiled.h"
#include "mlx/graph_utils.h"
#include "mlx/primitives.h" #include "mlx/primitives.h"
#include "mlx/utils.h" #include "mlx/utils.h"
@ -81,13 +82,27 @@ std::string build_lib_name(
const std::vector<array>& outputs, const std::vector<array>& outputs,
const std::vector<array>& tape, const std::vector<array>& tape,
const std::unordered_set<uintptr_t>& constant_ids) { const std::unordered_set<uintptr_t>& constant_ids) {
NodeNamer namer;
std::ostringstream os; std::ostringstream os;
std::ostringstream constant_hasher; std::ostringstream constant_hasher;
// Fill the input names. This is not really necessary, I just like having A,
// B, C, ... as the inputs.
for (auto& x : inputs) {
namer.get_name(x);
}
// The primitives describing the tape. For unary and binary primitives this // The primitives describing the tape. For unary and binary primitives this
// must be enough to describe the full computation. // must be enough to describe the full computation.
for (auto& a : tape) { for (auto& a : tape) {
// name and type of output
os << namer.get_name(a) << kindof(a.dtype()) << a.itemsize();
// computation performed
a.primitive().print(os); a.primitive().print(os);
// name of inputs to the function
for (auto& inp : a.inputs()) {
os << namer.get_name(inp);
}
} }
os << "_"; os << "_";