zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-26 10:41:14 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity
76c919b4ec
mlx/mlx/backend/no_metal/primitives.cpp

108 lines
2.0 KiB
C++
Raw Normal View History

Custom VJP and checkpointing (#541) * Implement custom_vjp and checkpointing * Add a dependency management primitive * Change the eval order to deep branches first * Add graph depth tracking to the array
2024-01-31 08:04:45 +08:00
// Copyright © 2023-2024 Apple Inc.
copyright + ack
2023-12-01 03:12:53 +08:00
jagrit's commit files
2023-11-30 02:52:08 +08:00
#include "mlx/primitives.h"
Separate fast ops and primitives (#699)
2024-02-17 11:16:39 +08:00
#include "mlx/fast_primitives.h"
jagrit's commit files
2023-11-30 02:52:08 +08:00
Multi output primitives (#330) * Multi-output primitives --------- Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
2024-01-09 08:39:08 +08:00
#define NO_GPU_MULTI(func) \
void func::eval_gpu( \
const std::vector<array>& inputs, std::vector<array>& outputs) { \
throw std::runtime_error(#func " has no GPU implementation."); \
}
jagrit's commit files
2023-11-30 02:52:08 +08:00
#define NO_GPU(func) \
void func::eval_gpu(const std::vector<array>& inputs, array& out) { \
throw std::runtime_error(#func " has no GPU implementation."); \
}
namespace mlx::core {
NO_GPU(Abs)
NO_GPU(Add)
Update GEMM (#424) * Organize and collect metal subroutine templates and elements in `metal/kernels/steel/` * Update gemm elements for better performance * Add split-K specialization for gemm * Add `addmm` primitive, op and bindings for fused matmul and bias addition * Update tests and benchmarks as needed
2024-01-18 04:42:39 +08:00
NO_GPU(AddMM)
jagrit's commit files
2023-11-30 02:52:08 +08:00
NO_GPU(Arange)
NO_GPU(ArcCos)
NO_GPU(ArcCosh)
NO_GPU(ArcSin)
NO_GPU(ArcSinh)
NO_GPU(ArcTan)
NO_GPU(ArcTanh)
NO_GPU(ArgPartition)
NO_GPU(ArgReduce)
NO_GPU(ArgSort)
NO_GPU(AsType)
NO_GPU(AsStrided)
NO_GPU(Broadcast)
Floor and Ceil (#150) * Implements Floor and Ceil Ops
2023-12-15 02:00:23 +08:00
NO_GPU(Ceil)
Compile primitive (#571) * Compiled primitive with basic binary, unary graph-level fusion
2024-02-05 22:51:22 +08:00
NO_GPU_MULTI(Compiled)
jagrit's commit files
2023-11-30 02:52:08 +08:00
NO_GPU(Concatenate)
NO_GPU(Convolution)
NO_GPU(Copy)
NO_GPU(Cos)
NO_GPU(Cosh)
Custom VJP and checkpointing (#541) * Implement custom_vjp and checkpointing * Add a dependency management primitive * Change the eval order to deep branches first * Add graph depth tracking to the array
2024-01-31 08:04:45 +08:00
NO_GPU_MULTI(CustomVJP)
NO_GPU_MULTI(Depends)
jagrit's commit files
2023-11-30 02:52:08 +08:00
NO_GPU(Divide)
Compile primitive (#571) * Compiled primitive with basic binary, unary graph-level fusion
2024-02-05 22:51:22 +08:00
NO_GPU_MULTI(DivMod)
NumberOfElements for shapeless compile and vmap fixes (#802)
2024-03-14 01:34:14 +08:00
NO_GPU(NumberOfElements)
Add the remainder op (#85) * Add remainder in the C++ backend * Add the python binding and test
2023-12-09 07:08:52 +08:00
NO_GPU(Remainder)
jagrit's commit files
2023-11-30 02:52:08 +08:00
NO_GPU(Equal)
NO_GPU(Erf)
NO_GPU(ErfInv)
NO_GPU(Exp)
NO_GPU(FFT)
Floor and Ceil (#150) * Implements Floor and Ceil Ops
2023-12-15 02:00:23 +08:00
NO_GPU(Floor)
jagrit's commit files
2023-11-30 02:52:08 +08:00
NO_GPU(Full)
NO_GPU(Gather)
NO_GPU(Greater)
NO_GPU(GreaterEqual)
NO_GPU(Less)
NO_GPU(LessEqual)
NO_GPU(Load)
NO_GPU(Log)
NO_GPU(Log1p)
NO_GPU(LogicalNot)
feat: add logicalAnd and logicalOR (#386) * feat: add logicalAnd and logicalOR * run pre-commit * Refactor logical_and and logical_or functions * Add acknowledgement * Add logical AND and logical OR operators * Refactor logical_and and logical_or functions * Add support for logical operators on bool arrays * Update mlx/ops.cpp Co-authored-by: Awni Hannun <awni.hannun@gmail.com> * Update mlx/ops.cpp Co-authored-by: Awni Hannun <awni.hannun@gmail.com> * Add logical AND and OR operators for arrays and scalars * Refactor vjp and jvp methods in primitives.cpp * Add overloaded operators for logical AND and OR * format --------- Co-authored-by: Awni Hannun <awni.hannun@gmail.com> Co-authored-by: Awni Hannun <awni@apple.com>
2024-01-08 23:00:05 +08:00
NO_GPU(LogicalAnd)
NO_GPU(LogicalOr)
jagrit's commit files
2023-11-30 02:52:08 +08:00
NO_GPU(LogAddExp)
NO_GPU(Matmul)
NO_GPU(Maximum)
NO_GPU(Minimum)
NO_GPU(Multiply)
NO_GPU(Negative)
NO_GPU(NotEqual)
NO_GPU(Pad)
NO_GPU(Partition)
NO_GPU(Power)
Compile primitive (#571) * Compiled primitive with basic binary, unary graph-level fusion
2024-02-05 22:51:22 +08:00
NO_GPU_MULTI(QRF)
An initial quantized matmul implementation (#205) * Add quantized matvec * Add quantized matrix matrix with 2nd matrix transposed * Add quantized matmul tests * Add a slow cpu quantized matmul * Add a slightly faster vectorized cpu version
2023-12-19 15:18:57 +08:00
NO_GPU(QuantizedMatmul)
jagrit's commit files
2023-11-30 02:52:08 +08:00
NO_GPU(RandomBits)
NO_GPU(Reduce)
NO_GPU(Reshape)
Adds round op and primitive (#203)
2023-12-19 03:32:48 +08:00
NO_GPU(Round)
jagrit's commit files
2023-11-30 02:52:08 +08:00
NO_GPU(Scan)
NO_GPU(Scatter)
Implement the 'where' primitive for conditional selection (#664)
2024-02-23 07:10:48 +08:00
NO_GPU(Select)
jagrit's commit files
2023-11-30 02:52:08 +08:00
NO_GPU(Sigmoid)
NO_GPU(Sign)
NO_GPU(Sin)
NO_GPU(Sinh)
NO_GPU(Slice)
NO_GPU(Softmax)
NO_GPU(Sort)
Split multi output (#461) * Multi-output split primitive * Add the multi-output split to the ArrayIterator * Add some grad tests for split
2024-01-17 05:33:55 +08:00
NO_GPU_MULTI(Split)
jagrit's commit files
2023-11-30 02:52:08 +08:00
NO_GPU(Square)
NO_GPU(Sqrt)
NO_GPU(StopGradient)
NO_GPU(Subtract)
Add SVD primitive (#809) Add SVD op using Accelerate's LAPACK following https://developer.apple.com/documentation/accelerate/ compressing_an_image_using_linear_algebra Co-authored-by: Nicolo Valigi <nvaligi@apple.com>
2024-03-13 03:30:11 +08:00
NO_GPU_MULTI(SVD)
jagrit's commit files
2023-11-30 02:52:08 +08:00
NO_GPU(Tan)
NO_GPU(Tanh)
NO_GPU(Transpose)
Compile primitive (#571) * Compiled primitive with basic binary, unary graph-level fusion
2024-02-05 22:51:22 +08:00
Custom primitive + RoPE fat op (#676) * extensions start * rope custom op * fix build * docs + rope benchmark * fix test * Add a Metal kernel for RoPE * Fix position of traditional * transform tests * Move rope computation to float and fix tests * Fix the test and a typo * change to fast * fix no metal build --------- Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
2024-02-15 06:04:25 +08:00
namespace fast {
NO_GPU_MULTI(RoPE)
Fast Inference SDPA op (#735) * Fast Inference SDPA op Implements metal shaders for: o = mx.fast_inference_sdpa(queries, keys, values, scale, mask) Supports fp16, fp32 dtypes; assumes d_k = 128. Generic op support / prompt encoding supported via mlx primitives. Metal implementation is for the inference use case only. Majority of performance benefits appears to results from GQA & reduced bandwidth requirements; there is approximate performance parity for the MHA use case (from some measurements on M3 Max). * Flush shared memory to zero before unprotected reads for (scores @ values) * Move to fast:: namespace, address reviewer comments ... also attempt to revert formatter auto-change for files not relevant to this change * Shared memory flush to top of kernel * Resolve compiler warnings * Update python/src/fast.cpp Co-authored-by: Awni Hannun <awni.hannun@gmail.com> * Update python/src/fast.cpp Co-authored-by: Awni Hannun <awni.hannun@gmail.com> * Update python/src/fast.cpp Co-authored-by: Awni Hannun <awni.hannun@gmail.com> * Update python/src/fast.cpp Co-authored-by: Awni Hannun <awni.hannun@gmail.com> * Update docstring per PR feedback * Softmax in higher precision, ... * route to fallback for more use cases - batch size > 1, head_dim other than 128, etc. * Address linux build failure * Address other reviewer comments * Remove extraneous eval_cpu function per review --------- Co-authored-by: Atila Orhon <64497909+atiorh@users.noreply.github.com> Co-authored-by: Awni Hannun <awni.hannun@gmail.com> Co-authored-by: atila <atiorh@icloud.com>
2024-03-05 13:06:11 +08:00
NO_GPU(ScaledDotProductAttention)
Custom primitive + RoPE fat op (#676) * extensions start * rope custom op * fix build * docs + rope benchmark * fix test * Add a Metal kernel for RoPE * Fix position of traditional * transform tests * Move rope computation to float and fix tests * Fix the test and a typo * change to fast * fix no metal build --------- Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
2024-02-15 06:04:25 +08:00
} // namespace fast
jagrit's commit files
2023-11-30 02:52:08 +08:00
} // namespace mlx::core
Reference in New Issue Copy Permalink