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

Feature complete Metal FFT (#1102)

Browse Source 
* feature complete metal fft

* fix contiguity bug

* jit fft

* simplify rader/bluestein constant computation

* remove kernel/utils.h dep

* remove bf16.h dep

* format

---------

Co-authored-by: Alex Barron <abarron22@apple.com>
This commit is contained in:
Alex Barron
2024-06-06 12:57:25 -07:00
committed by GitHub
parent 0e585b4409
commit 27d70c7d9d
17 changed files with 2601 additions and 367 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
mlx/backend/metal/nojit_kernels.cpp
View File

@@ -191,4 +191,17 @@ MTL::ComputePipelineState* get_steel_conv_general_kernel(
return d.get_kernel(kernel_name);
}
MTL::ComputePipelineState* get_fft_kernel(
metal::Device& d,
const std::string& kernel_name,
const std::string& hash_name,
const int tg_mem_size,
const std::string& in_type,
const std::string& out_type,
int step,
bool real,
const metal::MTLFCList& func_consts) {
return d.get_kernel(kernel_name, "mlx", hash_name, func_consts);
}
} // namespace mlx::core