forked from pytorch/FBGEMM
-
Notifications
You must be signed in to change notification settings - Fork 0
/
defs.bzl
91 lines (86 loc) · 3.47 KB
/
defs.bzl
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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def get_fbgemm_base_srcs():
return [
"src/GenerateI8Depthwise.cc",
"src/RefImplementations.cc",
"src/Utils.cc",
]
def get_fbgemm_generic_srcs(with_base = False):
return [
"src/EmbeddingSpMDM.cc",
"src/EmbeddingSpMDMNBit.cc",
"src/ExecuteKernel.cc",
"src/ExecuteKernelU8S8.cc",
"src/Fbgemm.cc",
"src/FbgemmBfloat16Convert.cc",
"src/FbgemmConv.cc",
"src/FbgemmFP16.cc",
"src/FbgemmFloat16Convert.cc",
"src/FbgemmI64.cc",
"src/FbgemmI8Spmdm.cc",
"src/GenerateKernelU8S8S32ACC16.cc",
"src/GenerateKernelU8S8S32ACC16Avx512.cc", # Acc16 AVX512 JIT code gen
"src/GenerateKernelU8S8S32ACC16Avx512VNNI.cc",
"src/GenerateKernelU8S8S32ACC32.cc",
"src/GenerateKernelU8S8S32ACC32Avx512.cc", # Acc32 AVX512 JIT code gen
"src/GenerateKernelU8S8S32ACC32Avx512VNNI.cc",
"src/GroupwiseConvAcc32Avx2.cc",
"src/PackAMatrix.cc",
"src/PackAWithIm2Col.cc",
"src/PackAWithQuantRowOffset.cc",
"src/PackAWithRowOffset.cc",
"src/PackBMatrix.cc",
"src/PackMatrix.cc",
"src/PackWeightMatrixForGConv.cc",
"src/PackWeightsForConv.cc",
"src/QuantUtils.cc",
"src/RowWiseSparseAdagradFused.cc",
"src/SparseAdagrad.cc",
"src/TransposeUtils.cc",
] + (get_fbgemm_base_srcs() if with_base else [])
def get_fbgemm_public_headers():
return [
"include/fbgemm/Fbgemm.h",
"include/fbgemm/FbgemmBuild.h",
"include/fbgemm/FbgemmFP16.h",
"include/fbgemm/FbgemmEmbedding.h",
"include/fbgemm/FbgemmConvert.h",
"include/fbgemm/FbgemmI64.h",
"include/fbgemm/FbgemmI8DepthwiseAvx2.h",
"include/fbgemm/OutputProcessing-inl.h",
"include/fbgemm/PackingTraits-inl.h",
"include/fbgemm/QuantUtils.h",
"include/fbgemm/QuantUtilsAvx2.h",
"include/fbgemm/Utils.h",
"include/fbgemm/UtilsAvx2.h",
"include/fbgemm/ConvUtils.h",
"include/fbgemm/Types.h",
"include/fbgemm/FbgemmI8Spmdm.h",
]
def get_fbgemm_avx2_srcs(msvc = False):
return [
#All the source files that either use avx2 instructions statically
"src/EmbeddingSpMDMAvx2.cc",
"src/FbgemmBfloat16ConvertAvx2.cc",
"src/FbgemmFloat16ConvertAvx2.cc",
"src/FbgemmI8Depthwise3DAvx2.cc",
"src/FbgemmI8DepthwiseAvx2.cc",
"src/FbgemmI8DepthwisePerChannelQuantAvx2.cc",
"src/OptimizedKernelsAvx2.cc",
"src/PackDepthwiseConvMatrixAvx2.cc",
"src/QuantUtilsAvx2.cc",
"src/UtilsAvx2.cc",
#FP16 kernels contain inline assembly and inline assembly syntax for MSVC is different.
"src/FbgemmFP16UKernelsAvx2.cc" if not msvc else "src/FbgemmFP16UKernelsIntrinsicAvx2.cc",
]
def get_fbgemm_avx512_srcs(msvc = False):
return [
#All the source files that use avx512 instructions statically
"src/FbgemmBfloat16ConvertAvx512.cc",
"src/FbgemmFloat16ConvertAvx512.cc",
"src/UtilsAvx512.cc",
#FP16 kernels contain inline assembly and inline assembly syntax for MSVC is different.
"src/FbgemmFP16UKernelsAvx512.cc" if not msvc else "src/FbgemmFP16UKernelsIntrinsicAvx512.cc",
"src/FbgemmFP16UKernelsAvx512_256.cc" if not msvc else "src/FbgemmFP16UKernelsIntrinsicAvx512_256.cc",
]
def get_fbgemm_tests(skip_tests = []):
return native.glob(["test/*Test.cc"], exclude = skip_tests)