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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
|
#pragma once
#include <ATen/code_template.h>
namespace torch {
namespace jit {
namespace fuser {
namespace cpu {
/*with type_as not checking type of its input, a fusion group can have non-fp32
tensor as input. Correct code for this case is generated, however, nvrtc does
not know how to handle int*_t integer types, so typedefs help it handle those
cases*/
static auto type_declarations_template = at::jit::CodeTemplate(R"(
#define POS_INFINITY INFINITY
#define NEG_INFINITY -INFINITY
typedef ${IndexType} IndexType;
template<typename T, size_t N>
struct TensorInfo {
T* data;
IndexType sizes[N];
IndexType strides[N];
};
template<typename T>
struct TensorInfo<T, 0> {
T * data;
};
)");
static auto cpu_compilation_unit_template = at::jit::CodeTemplate(R"(
#include <math.h>
#include <cstddef>
#include <cstdint>
double rsqrt(double x) {
return 1.0/sqrt(x);
}
float rsqrtf(float x) {
return 1.0f/sqrtf(x);
}
double frac(double x) {
return x - trunc(x);
}
float fracf(float x) {
return x - truncf(x);
}
${type_declarations}
#ifdef _MSC_VER
template<size_t n> struct int_of_size;
#define DEFINE_INT_OF_SIZE(int_t) \
template<> struct int_of_size<sizeof(int_t)> { using type = int_t; }
DEFINE_INT_OF_SIZE(int64_t);
DEFINE_INT_OF_SIZE(int32_t);
DEFINE_INT_OF_SIZE(int16_t);
DEFINE_INT_OF_SIZE(int8_t);
#undef DEFINE_INT_OF_SIZE
template <typename T>
using int_same_size_t = typename int_of_size<sizeof(T)>::type;
#define IndexTypeLoop int_same_size_t<IndexType>
#define ToIndexTypeLoop(x) static_cast<IndexTypeLoop>(x)
#else
#define IndexTypeLoop IndexType
#define ToIndexTypeLoop(x) x
#endif
#define OMP_THRESHOLD 100000
static void ${kernelName}_kernel(IndexType totalElements, ${formals}) {
#pragma omp parallel for if(totalElements > OMP_THRESHOLD)
for (IndexTypeLoop linearIndex = 0;
linearIndex < ToIndexTypeLoop(totalElements);
linearIndex += 1) {
// Convert `linearIndex` into an offset of tensor:
${tensorOffsets}
// calculate the results
${kernelBody}
}
}
#ifdef _WIN32
#define JIT_API __declspec(dllexport)
#else
#define JIT_API
#endif
extern "C"
JIT_API void ${kernelName}(IndexType totalElements, void ** args) {
${kernelName}_kernel(totalElements ${,argument_loads});
}
)");
} // namespace cpu
} // namespace fuser
} // namespace jit
} // namespace torch
|
