Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

add thread local cache for brgemm #350

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

add thread local cache for brgemm #350

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
29f25e2
add thread local cache for brgemm
crazydemo Sep 19, 2024
37267a7
encapsulate thread local cache
crazydemo Sep 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
103 changes: 79 additions & 24 deletions lib/gc/ExecutionEngine/CPURuntime/Microkernel/BrgemmOnednn.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,53 @@ static std::vector<brgemm_desc_t> g_brgemm_desc_list;
static std::vector<brgemm_kernel_t *> g_brgemm_kernel_list;
static std::vector<std::unique_ptr<char[]>> g_brgemm_palette;

struct brgemm_cache_info_t {
std::shared_ptr<brgemm_desc_t> desc;
std::shared_ptr<brgemm_kernel_t> kernel;
std::shared_ptr<char> palette;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we could use shared_ptr<char[]> for palette, and for desc and kernel we don't need to change, since they are not managed by smart ptr, and storing ptr to vector elements is dangerous as well

Copy link
Contributor

@yifeizh2 yifeizh2 Sep 20, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As we created brgemm_cache_info_t to store desc/kernel/palette together thread locally, would it be a better manner to also use brgemm_cache_info_t for global management?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it's a good idea, we can unify the struct used in both thead-local and global


brgemm_cache_info_t() = default;
brgemm_cache_info_t(brgemm_desc_t *d, brgemm_kernel_t *k, char *p)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ideally we need to change the unique_ptr in global palette pool to shared_ptr as well, and pass the shared_ptr of palette here for construction

: desc(d), kernel(k), palette(p) {}
brgemm_cache_info_t &operator=(const brgemm_cache_info_t &other) {
if (this != &other) {
desc = other.desc;
kernel = other.kernel;
palette = other.palette;
}
return *this;
}
};

class brgemm_cache_manager {
public:
static brgemm_cache_manager &getInstance() {
static thread_local brgemm_cache_manager instance;
return instance;
}

void insertOrUpdate(int64_t key, const brgemm_cache_info_t &info) {
cache_[key] = info;
}

bool tryGet(int64_t key, brgemm_cache_info_t &info) {
auto it = cache_.find(key);
if (it != cache_.end()) {
info = it->second;
return true;
}
return false;
}

private:
brgemm_cache_manager() {}

brgemm_cache_manager(const brgemm_cache_manager &) = delete;
brgemm_cache_manager &operator=(const brgemm_cache_manager &) = delete;

std::unordered_map<int64_t, brgemm_cache_info_t> cache_;
};

// TODO(haixin): use syscall to determine page size?
static constexpr size_t SCRATCH_SIZE = 2 * 4096;
// TODO(haixin): need to use custom thread management for scratch in the future?
Expand Down Expand Up @@ -105,25 +152,30 @@ int64_t dnnl_brgemm_dispatch(int64_t M, int64_t N, int64_t K, int64_t LDA,
g_brgemm_desc_list.push_back(desc);
g_brgemm_kernel_list.push_back(kernel);
g_brgemm_palette.emplace_back(palette_buffer);

return g_brgemm_desc_list.size() - 1;
}

void dnnl_brgemm_tileconfig(int64_t kernel_idx) {
char *palette_buffer = nullptr;
{
assert(kernel_idx >= 0 && "Invalid kernel handler");
auto &cache_manager = brgemm_cache_manager::getInstance();
brgemm_cache_info_t info;
if (!cache_manager.tryGet(kernel_idx, info)) {
read_lock_guard_t g(g_brgemm_lock);
assert(kernel_idx >= 0 && kernel_idx < (int64_t)g_brgemm_desc_list.size() &&
assert(kernel_idx < (int64_t)g_brgemm_desc_list.size() &&
"Invalid kernel handler");
brgemm_desc_t &desc = g_brgemm_desc_list[kernel_idx];
if (!desc.is_tmm) {
return;
}
palette_buffer = g_brgemm_palette[kernel_idx].get();
info = {&g_brgemm_desc_list[kernel_idx], g_brgemm_kernel_list[kernel_idx],
g_brgemm_palette[kernel_idx].get()};
cache_manager.insertOrUpdate(kernel_idx, info);
}
brgemm_desc_t *desc = info.desc.get();
std::shared_ptr<char> palette_buffer = info.palette;

assert(palette_buffer != nullptr && "Invalid palette for BRGEMM kernel");
amx_tile_configure(palette_buffer);
if (!desc->is_tmm) {
return;
}

assert(palette_buffer && "Invalid palette for BRGEMM kernel");
amx_tile_configure(palette_buffer.get());
}

void dnnl_brgemm_tilerelease() {
Expand All @@ -137,28 +189,31 @@ void dnnl_brgemm_tilerelease() {
void dnnl_brgemm_execute(int64_t kernel_idx, void *A, uint64_t A_offset,
void *B, uint64_t B_offset, void *C, uint64_t C_offset,
int num) {
brgemm_kernel_t *kernel = nullptr;
size_t A_offset_in_bytes;
size_t B_offset_in_bytes;
size_t C_offset_in_bytes;
{
auto &cache_manager = brgemm_cache_manager::getInstance();
brgemm_cache_info_t info;
if (!cache_manager.tryGet(kernel_idx, info)) {
read_lock_guard_t g(g_brgemm_lock);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since it's thread local, do we still need this lock?

assert(kernel_idx >= 0 && kernel_idx < (int64_t)g_brgemm_desc_list.size() &&
"Invalid kernel handler");
info = {&g_brgemm_desc_list[kernel_idx], g_brgemm_kernel_list[kernel_idx],
g_brgemm_palette[kernel_idx].get()};
cache_manager.insertOrUpdate(kernel_idx, info);
}

brgemm_desc_t &desc = g_brgemm_desc_list[kernel_idx];
kernel = g_brgemm_kernel_list[kernel_idx];
assert(info.kernel && "Invalid brgemm kernel pointer");

A_offset_in_bytes = dnnl::impl::types::data_type_size(desc.dt_a) * A_offset;
B_offset_in_bytes = dnnl::impl::types::data_type_size(desc.dt_b) * B_offset;
C_offset_in_bytes = dnnl::impl::types::data_type_size(desc.dt_c) * C_offset;
}
size_t A_offset_in_bytes =
dnnl::impl::types::data_type_size(info.desc->dt_a) * A_offset;
size_t B_offset_in_bytes =
dnnl::impl::types::data_type_size(info.desc->dt_b) * B_offset;
size_t C_offset_in_bytes =
dnnl::impl::types::data_type_size(info.desc->dt_c) * C_offset;

assert(kernel && "Invalid brgemm kernel pointer");
char *A_arith = (char *)A;
char *B_arith = (char *)B;
char *C_arith = (char *)C;
brgemm_kernel_execute(kernel, num, (void *)(A_arith + A_offset_in_bytes),
brgemm_kernel_execute(info.kernel.get(), num,
(void *)(A_arith + A_offset_in_bytes),
(void *)(B_arith + B_offset_in_bytes), nullptr,
(void *)(C_arith + C_offset_in_bytes), (void *)scratch);
}
Expand Down
Loading