Fix ggml_vk_h2d_tensor throwing on second call

This commit is contained in:
niansa 2023-06-23 16:50:37 +02:00
parent 09b0b3a49b
commit 98e588c6eb
2 changed files with 21 additions and 13 deletions

View file

@ -91,12 +91,20 @@ void ggml_vk_h2d_tensor(struct ggml_kompute_context * ctx, struct ggml_tensor *
auto data = t->data;
auto size = ggml_nbytes(t);
std::vector<byte> vec(size);
memcpy(vec.data(), data, size);
auto res = ctx->tensors.find(t);
auto tensor = mgr.tensorT<byte>(vec);
mgr.sequence()->eval<kp::OpTensorSyncDevice>({tensor});
ctx->tensors.emplace(t, std::move(tensor));
if (res != ctx->tensors.end()) {
assert(res->second->size() != size);
res->second->setRawData(data);
mgr.sequence()->eval<kp::OpTensorSyncDevice>({res->second});
} else {
std::vector<byte> vec(size);
memcpy(vec.data(), data, size);
auto tensor = mgr.tensorT<byte>(vec);
mgr.sequence()->eval<kp::OpTensorSyncDevice>({tensor});
ctx->tensors.emplace(t, std::move(tensor));
}
}
void ggml_vk_d2h_tensor(struct ggml_kompute_context * ctx, struct ggml_tensor * t) {

View file

@ -2769,7 +2769,7 @@ struct llama_context * llama_init_from_file(
}
#elif defined(GGML_USE_KOMPUTE)
if (params.n_gpu_layers > 0) {
// this allocates all Metal resources and memory buffers
// this allocates all Vulkan resources and memory buffers
ctx->ctx_kompute = ggml_vk_init();
void * data_ptr = NULL;
@ -2787,21 +2787,21 @@ struct llama_context * llama_init_from_file(
printf("%s: max tensor size = %8.2f MB\n", __func__, max_size/1024.0/1024.0);
#define LLAMA_METAL_CHECK_BUF(result) \
#define LLAMA_VK_CHECK_BUF(result) \
if (!(result)) { \
fprintf(stderr, "%s: failed to add buffer\n", __func__); \
llama_free(ctx); \
return NULL; \
}
LLAMA_METAL_CHECK_BUF(ggml_vk_add_buffer(ctx->ctx_kompute, "data", data_ptr, data_size, max_size));
LLAMA_VK_CHECK_BUF(ggml_vk_add_buffer(ctx->ctx_kompute, "data", data_ptr, data_size, max_size));
LLAMA_METAL_CHECK_BUF(ggml_vk_add_buffer(ctx->ctx_kompute, "eval", ctx->buf_compute.addr, ctx->buf_compute.size, 0));
LLAMA_METAL_CHECK_BUF(ggml_vk_add_buffer(ctx->ctx_kompute, "kv", ctx->model.kv_self.buf.addr, ctx->model.kv_self.buf.size, 0));
LLAMA_VK_CHECK_BUF(ggml_vk_add_buffer(ctx->ctx_kompute, "eval", ctx->buf_compute.addr, ctx->buf_compute.size, 0));
LLAMA_VK_CHECK_BUF(ggml_vk_add_buffer(ctx->ctx_kompute, "kv", ctx->model.kv_self.buf.addr, ctx->model.kv_self.buf.size, 0));
LLAMA_METAL_CHECK_BUF(ggml_vk_add_buffer(ctx->ctx_kompute, "scr0", ctx->buf_scratch[0].addr, ctx->buf_scratch[0].size, 0));
LLAMA_METAL_CHECK_BUF(ggml_vk_add_buffer(ctx->ctx_kompute, "scr1", ctx->buf_scratch[1].addr, ctx->buf_scratch[1].size, 0));
#undef LLAMA_METAL_CHECK_BUF
LLAMA_VK_CHECK_BUF(ggml_vk_add_buffer(ctx->ctx_kompute, "scr0", ctx->buf_scratch[0].addr, ctx->buf_scratch[0].size, 0));
LLAMA_VK_CHECK_BUF(ggml_vk_add_buffer(ctx->ctx_kompute, "scr1", ctx->buf_scratch[1].addr, ctx->buf_scratch[1].size, 0));
#undef LLAMA_VK_CHECK_BUF
}
#endif