@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
hpp 146 lines 3.21 KB
Raw
1 #ifndef UTILS_HPP
2 #define UTILS_HPP
3
4 #include <atomic>
5 #include <string>
6 #include <unistd.h>
7 #include <ATen/cpu/Utils.h>
8
9 #include "cpu/cpu_types.hpp"
10
11 namespace cpu_utils {
12 enum class ISA { AMX, VEC, RVV, NEON, VSX };
13
14 inline ISA get_isa(const std::string& isa) {
15 if (isa == "amx") {
16 return ISA::AMX;
17 } else if (isa == "vec") {
18 return ISA::VEC;
19 } else if (isa == "rvv") {
20 return ISA::RVV;
21 } else if (isa == "neon") {
22 return ISA::NEON;
23 } else if (isa == "vsx") {
24 return ISA::VSX;
25 } else {
26 TORCH_CHECK(false, "Invalid isa type: " + isa);
27 }
28 }
29
30 template <typename T>
31 struct VecTypeTrait {
32 using vec_t = void;
33 };
34
35 template <>
36 struct VecTypeTrait<float> {
37 using vec_t = vec_op::FP32Vec16;
38 };
39
40 template <>
41 struct VecTypeTrait<c10::BFloat16> {
42 using vec_t = vec_op::BF16Vec16;
43 };
44
45 template <>
46 struct VecTypeTrait<c10::Half> {
47 using vec_t = vec_op::FP16Vec16;
48 };
49
50 struct Counter {
51 std::atomic<int64_t> counter;
52 char _padding[56];
53
54 Counter() : counter(0) {}
55
56 void reset_counter() { counter.store(0); }
57
58 int64_t acquire_counter() { return counter++; }
59 };
60
61 inline int64_t get_available_l2_size() {
62 #if defined(__s390x__) || defined(__powerpc__)
63 static int64_t size = []() {
64 uint32_t l2_cache_size = 0;
65 auto caps = at::cpu::get_cpu_capabilities();
66 auto it = caps.find("l2_cache_size");
67 if (it != caps.end()) {
68 l2_cache_size = static_cast<uint32_t>(it->second.toInt());
69 }
70 if (l2_cache_size == 0) {
71 long sys_l2 = sysconf(_SC_LEVEL2_CACHE_SIZE);
72 if (sys_l2 > 0) {
73 l2_cache_size = static_cast<uint32_t>(sys_l2);
74 }
75 }
76 if (l2_cache_size == 0) {
77 l2_cache_size = 256 * 1024;
78 }
79 return static_cast<int64_t>(l2_cache_size) >> 1;
80 }();
81 return size;
82 #else
83 static int64_t size = []() {
84 auto caps = at::cpu::get_cpu_capabilities();
85 const uint32_t l2_cache_size = caps.at("l2_cache_size").toInt();
86 return l2_cache_size >> 1;
87 }();
88 return size;
89 #endif
90 }
91
92 template <int32_t alignment_v, typename T>
93 inline T round_up(T size) {
94 T alignment = alignment_v;
95 return (((size + alignment - 1) / alignment) * alignment);
96 }
97
98 template <int32_t alignment_v, typename T>
99 inline T round_down(T size) {
100 T alignment = alignment_v;
101 return (size / alignment) * alignment;
102 }
103
104 template <typename T>
105 inline void print_logits(const char* name, T* ptr, int32_t row, int32_t col,
106 int32_t stride) {
107 std::stringstream ss;
108 ss << std::fixed << std::setprecision(5) << name << ": [\n";
109 auto* curr_logits_buffer = ptr;
110 for (int32_t m = 0; m < row; ++m) {
111 for (int32_t n = 0; n < col; ++n) {
112 ss << curr_logits_buffer[n] << ", ";
113 }
114 ss << "\n";
115 curr_logits_buffer += stride;
116 }
117 ss << "]\n";
118 std::printf("%s", ss.str().c_str());
119 }
120
121 class ScratchPadManager {
122 public:
123 static constexpr size_t allocation_unit = 4 * 1024; // 4KB
124
125 static ScratchPadManager* get_scratchpad_manager();
126
127 ScratchPadManager();
128
129 template <typename T>
130 T* get_data() {
131 return reinterpret_cast<T*>(ptr_);
132 }
133
134 static size_t round(size_t size) {
135 return ((size + allocation_unit - 1) / allocation_unit) * allocation_unit;
136 }
137
138 void realloc(size_t new_size);
139
140 private:
141 size_t size_;
142 void* ptr_;
143 };
144 } // namespace cpu_utils
145
146 #endif