MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / gen

Function gen

dnn/src/x86/script/gen_convolution2.py:4–149  ·  view source on GitHub ↗
(mode, simd, fsize)

Source from the content-addressed store, hash-verified

2import itertools
3
4def gen(mode, simd, fsize):
5 funcname = "convolution_{mode}_fh{fsize}_{simd}".format(**vars())
6 filename = funcname + ".cpp"
7 if simd == 'fma':
8 MAX_H = 15 - fsize
9 elif simd == 'avx' or simd == 'sse':
10 MAX_H = 14 - fsize
11 else:
12 assert False
13 if simd == "sse":
14 width = 4
15 mm_type = "__m128"
16 mm_load = "_mm_loadu_ps"
17 mm_store = "_mm_storeu_ps"
18 mm_mul = "_mm_mul_ps"
19 mm_add = "_mm_add_ps"
20 mm_set1 = "_mm_set1_ps"
21 mm_set0 = "_mm_setzero_ps"
22 mm_max = "_mm_max_ps"
23 mm_set1_sign = ""
24 header = ["xmmintrin.h"]
25 elif simd == "avx":
26 width = 8
27 mm_type = "__m256"
28 mm_load = "_mm256_loadu_ps"
29 mm_store = "_mm256_storeu_ps"
30 mm_mul = "_mm256_mul_ps"
31 mm_add = "_mm256_add_ps"
32 mm_set1 = "_mm256_broadcast_ss"
33 mm_set0 = "_mm256_setzero_ps"
34 mm_max = "_mm256_max_ps"
35 mm_set1_sign = "&"
36 header = ["immintrin.h", "avxintrin.h"]
37 elif simd == "fma":
38 width = 8
39 mm_type = "__m256"
40 mm_load = "_mm256_loadu_ps"
41 mm_store = "_mm256_storeu_ps"
42 mm_set1 = "_mm256_broadcast_ss"
43 mm_set0 = "_mm256_setzero_ps"
44 mm_max = "_mm256_max_ps"
45 mm_set1_sign = "&"
46 header = ["immintrin.h", "avxintrin.h", "fmaintrin.h"]
47 with open(filename, 'w') as f:
48 for H in range(1, MAX_H+1):
49 f.write("""#define SIMD_H{H} do {{ \\
50const size_t sh = dh; \\
51const float *src_d = src + sh*src_w; \\
52float *dst_d = dst + dh*dst_w; \\
53size_t dw = dst_w_beg; \\
54for (; dw < dst_w_end; dw += {width}) {{ \\
55 const size_t sw = dw; \\
56 float *dst_dd = dst_d + dw; \\
57 {mm_type} tmp0; \\
58""".format(**vars()))
59 if simd != "fma":
60 f.write(" {mm_type} tmp1; \\\n".format(**vars()))
61 for h in range(H):

Callers 15

RNGxorshfMethod · 0.50
create_graph_dumpFunction · 0.50
check_dump_by_computeFunction · 0.50
TESTFunction · 0.50
test_serializer_APlusBFunction · 0.50

Calls 2

formatMethod · 0.45
writeMethod · 0.45