@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
hpp 135 lines 4.67 KB
Raw
1 #ifndef CPU_MICRO_GEMM_IMPL_HPP
2 #define CPU_MICRO_GEMM_IMPL_HPP
3 #include "cpu/utils.hpp"
4 #include "cpu/cpu_types.hpp"
5
6 namespace cpu_micro_gemm {
7 #define DEFINE_CPU_MICRO_GEMM_PARAMS \
8 scalar_t *__restrict__ a_ptr, scalar_t *__restrict__ b_ptr, \
9 float *__restrict__ c_ptr, const int32_t m, const int32_t k, \
10 const int64_t lda, const int64_t b_n_group_stride, const int64_t ldc, \
11 const bool accum_c
12
13 #define CPU_MICRO_GEMM_PARAMS \
14 a_ptr, b_ptr, c_ptr, m, k, lda, b_n_group_stride, ldc, accum_c
15
16 // Note: weights for MicroGemm should be packed as (output_size / 16) contiguous
17 // blocks, means the logical shape of blocks is [16, input_size]. And the actual
18 // layout of blocks can be ISA-specific.
19 template <cpu_utils::ISA isa, typename scalar_t>
20 class MicroGemm {
21 public:
22 static constexpr int32_t MaxMSize = 16;
23 static constexpr int32_t NSize = 16;
24 static constexpr int32_t WeightOCGroupSize = 16;
25 // callers must pack A matrix before GEMM
26 static constexpr bool PackA = false;
27
28 public:
29 void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
30 TORCH_CHECK(false, "Unimplemented MicroGemm.");
31 }
32 };
33
34 template <cpu_utils::ISA isa, typename scalar_t>
35 class MicroGemmINT8;
36
37 template <int32_t n_size, typename scalar_t>
38 FORCE_INLINE void default_epilogue(float* __restrict__ c_ptr,
39 scalar_t* __restrict__ d_ptr,
40 const int32_t m, const int64_t ldc,
41 const int64_t ldd) {
42 using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
43 static_assert(n_size % 16 == 0);
44
45 float* __restrict__ curr_c = c_ptr;
46 scalar_t* __restrict__ curr_d = d_ptr;
47 for (int32_t i = 0; i < m; ++i) {
48 float* __restrict__ curr_c_iter = curr_c;
49 scalar_t* __restrict__ curr_d_iter = curr_d;
50 vec_op::unroll_loop<int32_t, n_size / 16>([&](int32_t n_g_idx) {
51 vec_op::FP32Vec16 c_vec_fp32(curr_c_iter);
52 scalar_vec_t c_vec(c_vec_fp32);
53 c_vec.save(curr_d_iter);
54 curr_c_iter += 16;
55 curr_d_iter += 16;
56 });
57 curr_c += ldc;
58 curr_d += ldd;
59 }
60 }
61
62 template <int32_t n_size, typename scalar_t>
63 FORCE_INLINE void bias_epilogue(float* __restrict__ c_ptr,
64 scalar_t* __restrict__ d_ptr,
65 scalar_t* __restrict__ bias_ptr,
66 const int32_t m, const int64_t ldc,
67 const int64_t ldd) {
68 using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
69 static_assert(n_size % 16 == 0);
70 constexpr int32_t n_group_num = n_size / 16;
71 static_assert(n_group_num <= 16);
72
73 vec_op::FP32Vec16 bias_vecs[n_group_num];
74 scalar_t* __restrict__ curr_bias = bias_ptr;
75 vec_op::unroll_loop<int32_t, n_group_num>([&](int32_t i) {
76 scalar_vec_t vec(curr_bias);
77 bias_vecs[i] = vec_op::FP32Vec16(vec);
78 curr_bias += 16;
79 });
80
81 float* __restrict__ curr_c = c_ptr;
82 scalar_t* __restrict__ curr_d = d_ptr;
83 for (int32_t i = 0; i < m; ++i) {
84 float* __restrict__ curr_c_iter = curr_c;
85 scalar_t* __restrict__ curr_d_iter = curr_d;
86 vec_op::unroll_loop<int32_t, n_group_num>([&](int32_t n_g_idx) {
87 vec_op::FP32Vec16 c_vec_fp32(curr_c_iter);
88 c_vec_fp32 = c_vec_fp32 + bias_vecs[n_g_idx];
89 scalar_vec_t c_vec(c_vec_fp32);
90 c_vec.save(curr_d_iter);
91 curr_c_iter += 16;
92 curr_d_iter += 16;
93 });
94 curr_c += ldc;
95 curr_d += ldd;
96 }
97 }
98
99 template <int32_t n_size, typename scalar_t>
100 FORCE_INLINE void add_bias_epilogue(float* c_ptr, float* d_ptr,
101 scalar_t* __restrict__ bias_ptr,
102 const int32_t m, const int64_t ldc,
103 const int64_t ldd) {
104 using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
105 static_assert(n_size % 16 == 0);
106 constexpr int32_t n_group_num = n_size / 16;
107 static_assert(n_group_num <= 16);
108
109 vec_op::FP32Vec16 bias_vecs[n_group_num];
110 scalar_t* __restrict__ curr_bias = bias_ptr;
111 vec_op::unroll_loop<int32_t, n_group_num>([&](int32_t i) {
112 scalar_vec_t vec(curr_bias);
113 bias_vecs[i] = vec_op::FP32Vec16(vec);
114 curr_bias += 16;
115 });
116
117 float* curr_c = c_ptr;
118 float* curr_d = d_ptr;
119 for (int32_t i = 0; i < m; ++i) {
120 float* curr_c_iter = curr_c;
121 float* curr_d_iter = curr_d;
122 vec_op::unroll_loop<int32_t, n_group_num>([&](int32_t n_g_idx) {
123 vec_op::FP32Vec16 c_vec_fp32(curr_c_iter);
124 c_vec_fp32 = c_vec_fp32 + bias_vecs[n_g_idx];
125 c_vec_fp32.save(curr_d_iter);
126 curr_c_iter += 16;
127 curr_d_iter += 16;
128 });
129 curr_c += ldc;
130 curr_d += ldd;
131 }
132 }
133 } // namespace cpu_micro_gemm
134
135 #endif