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

shapeless slice update and broadcast when possible (#1727)

Browse Source
This commit is contained in:
Awni Hannun
2024-12-23 11:25:15 -08:00
committed by GitHub
parent 0308e9af71
commit ebfe64b92d
6 changed files with 43 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
mlx/compile.cpp
View File

@@ -72,23 +72,6 @@ bool is_fusable(const Primitive& p) {
is_noop(p);
}
bool allows_shapeless(const Primitive& p) {
return typeid(p) == typeid(Arange) || typeid(p) == typeid(Compiled) ||
is_unary(p) || is_binary(p) || is_noop(p) || is_reduction(p) ||
typeid(p) == typeid(Softmax) || typeid(p) == typeid(Sort) ||
typeid(p) == typeid(ArgSort) || typeid(p) == typeid(ArgPartition) ||
typeid(p) == typeid(Partition) || typeid(p) == typeid(Select) ||
typeid(p) == typeid(NumberOfElements) || typeid(p) == typeid(Gather) ||
typeid(p) == typeid(Transpose) || typeid(p) == typeid(Concatenate) ||
typeid(p) == typeid(Matmul) || typeid(p) == typeid(QuantizedMatmul) ||
typeid(p) == typeid(Squeeze) || typeid(p) == typeid(ExpandDims) ||
typeid(p) == typeid(Flatten) || typeid(p) == typeid(Unflatten) ||
typeid(p) == typeid(fast::AffineQuantize) ||
typeid(p) == typeid(fast::LayerNorm) ||
typeid(p) == typeid(fast::RMSNorm) || typeid(p) == typeid(fast::RoPE) ||
typeid(p) == typeid(fast::ScaledDotProductAttention);
}
Compiled::Compiled(
Stream stream,
std::vector<array> inputs,
@@ -800,24 +783,6 @@ std::vector<array> compile_replace(
return outputs;
}
void compile_validate_shapeless(const std::vector<array>& tape) {
for (auto& t : tape) {
if (!t.has_primitive()) {
continue;
}
auto& p = t.primitive();
if (allows_shapeless(p)) {
continue;
}
std::ostringstream msg;
msg << "[compile] Cannot compile primitive ";
p.print(msg);
msg << " with shapeless enabled.";
throw std::invalid_argument(msg.str());
}
}
bool skip_compile() {
return compile_mode() == CompileMode::disabled ||
!(compile_available_for_device(default_device()));
@@ -877,10 +842,6 @@ std::function<std::vector<array>(const std::vector<array>&)> compile(
if (compile_mode() != CompileMode::no_fuse) {
compile_fuse(entry.tape, parents_map, entry.inputs, entry.outputs);
}
if (shapeless) {
compile_validate_shapeless(entry.tape);
}
}
// At this point we must have a tape, now replace the placeholders