MLX
Loading...
Searching...
No Matches
mma.h File Reference
#include <metal_simdgroup>
#include <metal_simdgroup_matrix>
#include <metal_stdlib>
#include "mlx/backend/metal/kernels/steel/defines.h"
#include "mlx/backend/metal/kernels/steel/gemm/transforms.h"
#include "mlx/backend/metal/kernels/steel/utils/integral_constant.h"

Go to the source code of this file.

Classes

struct  mlx::steel::BaseMMAFrag< T, kFragRows_, kFragCols_ >
 
struct  mlx::steel::BaseMMAFrag< T, 8, 8 >
 
struct  mlx::steel::MMATile< T, kTileRows_, kTileCols_, MMAFrag_ >
 
struct  mlx::steel::BlockMMA< T, U, BM, BN, BK, WM, WN, transpose_a, transpose_b, lda_tgp, ldb_tgp, AccumType, Epilogue >
 

Namespaces

namespace  mlx
 
namespace  mlx::steel
 

Functions

template<typename T , typename U , int M, int N, int K>
METAL_FUNC void mlx::steel::tile_matmad (thread MMATile< T, M, N > &D, thread MMATile< U, M, K > &A, thread MMATile< U, K, N > &B, thread MMATile< T, M, N > &C)