ggml-backend-meta: add multi-segment read support in get_tensor (#22063)
This commit is contained in:
@@ -1270,7 +1270,45 @@ static void ggml_backend_meta_buffer_get_tensor(ggml_backend_buffer_t buffer, co
|
|||||||
GGML_ASSERT(ggml_is_contiguous(tensor));
|
GGML_ASSERT(ggml_is_contiguous(tensor));
|
||||||
|
|
||||||
const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false);
|
const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false);
|
||||||
GGML_ASSERT(split_state.n_segments == 1);
|
|
||||||
|
if (split_state.n_segments != 1) {
|
||||||
|
GGML_ASSERT(split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS);
|
||||||
|
GGML_ASSERT(offset == 0);
|
||||||
|
GGML_ASSERT(size == ggml_nbytes(tensor));
|
||||||
|
GGML_ASSERT(tensor->ne[3] == 1);
|
||||||
|
size_t offset_data = 0;
|
||||||
|
std::vector<size_t> simple_offsets(n_bufs, 0);
|
||||||
|
if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_0) {
|
||||||
|
GGML_ASSERT(tensor->ne[2] == 1);
|
||||||
|
const int64_t blck_size = ggml_blck_size(tensor->type);
|
||||||
|
for (size_t s = 0; s < split_state.n_segments; s++) {
|
||||||
|
for (size_t j = 0; j < n_bufs; j++) {
|
||||||
|
const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j);
|
||||||
|
GGML_ASSERT(split_state.ne[s*n_bufs + j] % blck_size == 0);
|
||||||
|
const size_t nbytes = split_state.ne[s*n_bufs + j]/blck_size * tensor->nb[0];
|
||||||
|
ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_data, simple_offsets[j], nbytes,
|
||||||
|
tensor->ne[1], simple_tensor->nb[1], tensor->nb[1]);
|
||||||
|
offset_data += nbytes;
|
||||||
|
simple_offsets[j] += nbytes;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
GGML_ASSERT(offset_data*tensor->ne[1] == size);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_AXIS_1);
|
||||||
|
for (size_t s = 0; s < split_state.n_segments; s++) {
|
||||||
|
for (size_t j = 0; j < n_bufs; j++) {
|
||||||
|
const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j);
|
||||||
|
const size_t nbytes = split_state.ne[s*n_bufs + j] * tensor->nb[1];
|
||||||
|
ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_data, simple_offsets[j], nbytes,
|
||||||
|
tensor->ne[2], simple_tensor->nb[2], tensor->nb[2]);
|
||||||
|
offset_data += nbytes;
|
||||||
|
simple_offsets[j] += nbytes;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
GGML_ASSERT(offset_data*tensor->ne[2] == size);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
switch (split_state.axis) {
|
switch (split_state.axis) {
|
||||||
case GGML_BACKEND_SPLIT_AXIS_0:
|
case GGML_BACKEND_SPLIT_AXIS_0:
|
||||||
|
|||||||
Reference in New Issue
Block a user