forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
LinearAlgebra.h
32 lines (26 loc) · 1.15 KB
/
LinearAlgebra.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
#pragma once
#include <ATen/native/TransposeType.h>
#include <c10/util/complex.h>
namespace at {
namespace native {
void mkl_gemm_batched(
TransposeType trans_A, TransposeType trans_B,
int batch_size, int M, int N, int K, float alpha,
const float** A, int lda, const float** B, int ldb, float beta,
float** C, int ldc);
void mkl_gemm_batched(
TransposeType trans_A, TransposeType trans_B,
int batch_size, int M, int N, int K, double alpha,
const double** A, int lda, const double** B, int ldb, double beta,
double** C, int ldc);
void mkl_gemm_batched(
TransposeType trans_A, TransposeType trans_B,
int batch_size, int M, int N, int K, c10::complex<float> alpha,
const c10::complex<float>** A, int lda, const c10::complex<float>** B, int ldb,
c10::complex<float> beta, c10::complex<float>** C, int ldc);
void mkl_gemm_batched(
TransposeType trans_A, TransposeType trans_B,
int batch_size, int M, int N, int K, c10::complex<double> alpha,
const c10::complex<double>** A, int lda, const c10::complex<double>** B, int ldb,
c10::complex<double> beta, c10::complex<double>** C, int ldc);
}} // namespace at::native