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

Fp64 on the CPU (#1843)

Browse Source 
* add fp64 data type

* clean build

* update docs

* fix bug
This commit is contained in:
Awni Hannun
2025-02-07 15:52:22 -08:00
committed by GitHub
parent 1a1b2108ec
commit 1c0c118f7c
32 changed files with 438 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/distributed/mpi/mpi.cpp
View File

@@ -95,6 +95,7 @@ struct MPIWrapper {
LOAD_SYMBOL(ompi_mpi_int64_t, mpi_int64_);
LOAD_SYMBOL(ompi_mpi_uint64_t, mpi_uint64_);
LOAD_SYMBOL(ompi_mpi_float, mpi_float_);
LOAD_SYMBOL(ompi_mpi_double, mpi_double_);
LOAD_SYMBOL(ompi_mpi_c_complex, mpi_complex_);
}
@@ -164,6 +165,8 @@ struct MPIWrapper {
return mpi_float16_;
case bfloat16:
return mpi_bfloat16_;
case float64:
return mpi_double_;
}
}
@@ -218,6 +221,7 @@ struct MPIWrapper {
MPI_Datatype mpi_int64_;
MPI_Datatype mpi_uint64_;
MPI_Datatype mpi_float_;
MPI_Datatype mpi_double_;
MPI_Datatype mpi_complex_;
MPI_Datatype mpi_float16_;
MPI_Datatype mpi_bfloat16_;

4
mlx/distributed/ring/ring.cpp
View File

@@ -68,6 +68,10 @@
using T = float; \
__VA_ARGS__; \
} break; \
case float64: { \
using T = double; \
__VA_ARGS__; \
} break; \
case complex64: { \
using T = complex64_t; \
__VA_ARGS__; \