jeans/ollama
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ggml-backend: Ensure allocation meet backend requirements

Browse Source 
Backends can impose additional alignment requirements on buffer sizes.
We should ensure that we meet these or allocations can fail.
This commit is contained in:
Jesse Gross 2025-03-07 17:20:54 -08:00 committed by Jesse Gross
parent 6da8b6a879
commit 25f9b152f9
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
ml/backend/ggml/ggml.go
View File

@ -520,6 +520,10 @@ func shapeToGGML(shape []int) *C.int64_t {
return &sh[0] return &sh[0]
} }
func pad(length, pad C.size_t) C.size_t {
return ((length + pad - 1) / pad) * pad
}
func (c Context) newTensor(dtype ml.DType, shape []int) ml.Tensor { func (c Context) newTensor(dtype ml.DType, shape []int) ml.Tensor {
if c.buft == nil { if c.buft == nil {
panic("set Input, Output, or Layer before creating tensors") panic("set Input, Output, or Layer before creating tensors")
@ -551,7 +555,8 @@ func (c Context) newTensor(dtype ml.DType, shape []int) ml.Tensor {
} }
t := C.ggml_new_tensor(c.ctx, cdtype, C.int(len(shape)), shapeToGGML(shape)) t := C.ggml_new_tensor(c.ctx, cdtype, C.int(len(shape)), shapeToGGML(shape))
b := C.ggml_backend_buft_alloc_buffer(c.buft, C.ggml_nbytes(t)) size := pad(C.ggml_backend_buft_get_alloc_size(c.buft, t), C.ggml_backend_buft_get_alignment(c.buft))
b := C.ggml_backend_buft_alloc_buffer(c.buft, size)
C.ggml_backend_tensor_alloc(b, t, C.ggml_backend_buffer_get_base(b)) C.ggml_backend_tensor_alloc(b, t, C.ggml_backend_buffer_get_base(b))
return &Tensor{b: c.b, t: t} return &Tensor{b: c.b, t: t}
} }