@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
hpp 135 lines 3.94 KB
Raw
1 #ifndef CPU_MICRO_GEMM_VEC_HPP
2 #define CPU_MICRO_GEMM_VEC_HPP
3 #include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
4
5 namespace cpu_micro_gemm {
6 namespace {
7 // 8-2-16 pattern, 8 regs for A, 2 regs for B, 16 regs for C, [8, K] @ [k, 32]
8 template <typename scalar_t>
9 class TileGemm82 {
10 public:
11 FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
12 switch (m) {
13 case 1:
14 gemm_micro<1>(CPU_MICRO_GEMM_PARAMS);
15 break;
16 case 2:
17 gemm_micro<2>(CPU_MICRO_GEMM_PARAMS);
18 break;
19 case 3:
20 gemm_micro<3>(CPU_MICRO_GEMM_PARAMS);
21 break;
22 case 4:
23 gemm_micro<4>(CPU_MICRO_GEMM_PARAMS);
24 break;
25 case 5:
26 gemm_micro<5>(CPU_MICRO_GEMM_PARAMS);
27 break;
28 case 6:
29 gemm_micro<6>(CPU_MICRO_GEMM_PARAMS);
30 break;
31 case 7:
32 gemm_micro<7>(CPU_MICRO_GEMM_PARAMS);
33 break;
34 case 8:
35 gemm_micro<8>(CPU_MICRO_GEMM_PARAMS);
36 break;
37 }
38 }
39
40 template <int32_t M>
41 static void gemm_micro(DEFINE_CPU_MICRO_GEMM_PARAMS) {
42 static_assert(0 < M && M <= 8);
43 using load_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
44
45 scalar_t* __restrict__ curr_b_0 = b_ptr;
46 scalar_t* __restrict__ curr_b_1 = b_ptr + b_n_group_stride;
47 float* __restrict__ curr_c_0 = c_ptr;
48 float* __restrict__ curr_c_1 = c_ptr + 16;
49
50 vec_op::FP32Vec16 c_regs[M * 2];
51 if (accum_c) {
52 float* __restrict__ curr_m_c_0 = curr_c_0;
53 float* __restrict__ curr_m_c_1 = curr_c_1;
54 vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
55 c_regs[i * 2] = vec_op::FP32Vec16(curr_m_c_0);
56 c_regs[i * 2 + 1] = vec_op::FP32Vec16(curr_m_c_1);
57
58 // update
59 curr_m_c_0 += ldc;
60 curr_m_c_1 += ldc;
61 });
62 }
63
64 scalar_t* __restrict__ curr_a = a_ptr;
65 for (int32_t k_idx = 0; k_idx < k; ++k_idx) {
66 load_vec_t b_0_reg(curr_b_0);
67 vec_op::FP32Vec16 fp32_b_0_reg(b_0_reg);
68 load_vec_t b_1_reg(curr_b_1);
69 vec_op::FP32Vec16 fp32_b_1_reg(b_1_reg);
70
71 scalar_t* __restrict__ curr_m_a = curr_a;
72 vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
73 scalar_t v = *curr_m_a;
74 vec_op::FP32Vec16 a_reg(static_cast<float>(v));
75 c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg;
76 c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg;
77
78 // update
79 curr_m_a += lda;
80 });
81
82 // update
83 curr_a += 1;
84 curr_b_0 += 16;
85 curr_b_1 += 16;
86 }
87
88 vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
89 c_regs[i * 2].save(curr_c_0);
90 c_regs[i * 2 + 1].save(curr_c_1);
91
92 // update
93 curr_c_0 += ldc;
94 curr_c_1 += ldc;
95 });
96 }
97 };
98 } // namespace
99
100 // Gemm kernel uses vector instructions, requires B matrix to be packed
101 template <typename scalar_t>
102 class MicroGemm<cpu_utils::ISA::VEC, scalar_t> {
103 public:
104 static constexpr int32_t MaxMSize = 8;
105 static constexpr int32_t NSize = 32;
106 static constexpr int32_t WeightOCGroupSize = 16;
107 static constexpr bool PackA = false;
108
109 public:
110 void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
111 TileGemm82<scalar_t>::gemm(CPU_MICRO_GEMM_PARAMS);
112 }
113
114 // Note: pack contiguous weight [output_size, input_size] as contiguous
115 // packed weight [output_size / 16, input_size, 16]
116 static void pack_weight(const scalar_t* __restrict__ weight,
117 scalar_t* __restrict__ packed_weight,
118 const int32_t output_size, const int32_t input_size) {
119 TORCH_CHECK_EQ(output_size % 16, 0);
120 for (int32_t o_idx = 0; o_idx < output_size; ++o_idx) {
121 const scalar_t* __restrict__ curr_weight = weight + o_idx * input_size;
122 scalar_t* __restrict__ curr_packed_weight =
123 packed_weight + (o_idx / 16) * (16 * input_size) + o_idx % 16;
124 for (int32_t i_idx = 0; i_idx < input_size; ++i_idx) {
125 *curr_packed_weight = *curr_weight;
126
127 curr_packed_weight += 16;
128 ++curr_weight;
129 }
130 }
131 }
132 };
133 } // namespace cpu_micro_gemm
134
135 #endif