ggml.go 21.7 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
package ggml

3
4
5
6
7
8
// #cgo CPPFLAGS: -I${SRCDIR}/ggml/include
// #include <stdlib.h>
// #include <stdint.h>
// #include "ggml.h"
// #include "ggml-cpu.h"
// #include "ggml-backend.h"
Michael Yang's avatar
Michael Yang committed
9
10
11
import "C"

import (
12
	"errors"
Michael Yang's avatar
Michael Yang committed
13
14
	"fmt"
	"io"
15
	"iter"
Michael Yang's avatar
Michael Yang committed
16
	"log/slog"
17
	"maps"
Michael Yang's avatar
Michael Yang committed
18
	"os"
19
20
21
22
	"slices"
	"strconv"
	"strings"
	"unicode"
Michael Yang's avatar
Michael Yang committed
23
24
25
26
27
	"unsafe"

	"github.com/ollama/ollama/format"
	fs "github.com/ollama/ollama/fs/ggml"
	"github.com/ollama/ollama/ml"
28
	ggml "github.com/ollama/ollama/ml/backend/ggml/ggml/src"
Michael Yang's avatar
Michael Yang committed
29
30
31
	"golang.org/x/sync/errgroup"
)

32
33
func devices() iter.Seq[*C.struct_ggml_backend_device] {
	return func(yield func(*C.struct_ggml_backend_device) bool) {
34
		ggml.OnceLoad()
35
36
37
38
39
		for i := range C.ggml_backend_dev_count() {
			if !yield(C.ggml_backend_dev_get(i)) {
				return
			}
		}
Michael Yang's avatar
Michael Yang committed
40
	}
41
}
Michael Yang's avatar
Michael Yang committed
42
43

type Backend struct {
44
45
46
47
48
49
	meta    *fs.GGML
	sched   *C.struct_ggml_backend_sched
	tensors map[string]*C.struct_ggml_tensor
	input   *C.struct_ggml_backend
	output  *C.struct_ggml_backend
	layers  map[int]*C.struct_ggml_backend
50

51
	flashAttention bool
Michael Yang's avatar
Michael Yang committed
52
53
}

54
func New(r *os.File, params ml.BackendParams) (ml.Backend, error) {
Michael Yang's avatar
Michael Yang committed
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
	meta, n, err := fs.Decode(r, -1)
	if err != nil {
		return nil, err
	}

	slog.Info(
		"",
		"architecture", meta.KV().Architecture(),
		"file_type", meta.KV().FileType(),
		"name", meta.KV().String("general.name"),
		"description", meta.KV().String("general.description"),
		"num_tensors", len(meta.Tensors().Items()),
		"num_key_values", len(meta.KV()),
	)

70
	type deviceBufferType struct {
71
72
73
74
75
76
77
78
79
80
81
		d   *C.struct_ggml_backend_device
		bts []*C.struct_ggml_backend_buffer_type
	}

	var cpus, accels, gpus []*C.struct_ggml_backend_device
	for d := range devices() {
		switch C.ggml_backend_dev_type(d) {
		case C.GGML_BACKEND_DEVICE_TYPE_CPU:
			cpus = append(cpus, d)
		case C.GGML_BACKEND_DEVICE_TYPE_ACCEL:
			accels = append(accels, d)
Michael Yang's avatar
Michael Yang committed
82
		case C.GGML_BACKEND_DEVICE_TYPE_GPU:
83
			gpus = append(gpus, d)
Michael Yang's avatar
Michael Yang committed
84
85
86
		}
	}

Michael Yang's avatar
Michael Yang committed
87
	cpuDeviceBufferType := deviceBufferType{d: C.ggml_backend_dev_by_type(C.GGML_BACKEND_DEVICE_TYPE_CPU)}
88
89
90
91
	for _, d := range append(accels, append(gpus, cpus...)...) {
		switch C.ggml_backend_dev_type(d) {
		case C.GGML_BACKEND_DEVICE_TYPE_CPU,
			C.GGML_BACKEND_DEVICE_TYPE_ACCEL:
Michael Yang's avatar
Michael Yang committed
92
			cpuDeviceBufferType.bts = append(cpuDeviceBufferType.bts, C.ggml_backend_dev_buffer_type(d))
Michael Yang's avatar
Michael Yang committed
93
		}
94
95
	}

96
	var gpuDeviceBufferTypes []deviceBufferType
97
98
	for _, d := range gpus {
		bt := C.ggml_backend_dev_buffer_type(d)
99
		gpuDeviceBufferTypes = append(gpuDeviceBufferTypes, deviceBufferType{
100
			d:   d,
Michael Yang's avatar
Michael Yang committed
101
			bts: append([]*C.struct_ggml_backend_buffer_type{bt}, cpuDeviceBufferType.bts...),
102
		})
Michael Yang's avatar
Michael Yang committed
103
104
	}

105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
	splits := make([]float32, len(gpus))
	if func() bool {
		for _, s := range params.TensorSplit {
			if s != 0 {
				return true
			}
		}

		return false
	}() {
		splits = params.TensorSplit
	} else {
		for i := range splits {
			var free, total C.size_t
			C.ggml_backend_dev_memory(gpus[i], &free, &total)
			splits[i] = float32(free)
		}
	}

	var sum float32
	for i := range splits {
		sum += splits[i]
		splits[i] = sum
	}

130
	for i := range splits {
131
		splits[i] /= sum
132
133
	}

Michael Yang's avatar
Michael Yang committed
134
	input := cpuDeviceBufferType
135

136
	blocks := int(meta.KV().BlockCount())
Michael Yang's avatar
Michael Yang committed
137
	assignLayer := func(i int) deviceBufferType {
138
		if i >= params.NumGPULayers {
Michael Yang's avatar
Michael Yang committed
139
			return cpuDeviceBufferType
140
		}
141

142
143
		index := slices.IndexFunc(splits, func(f float32) bool { return float32(i)/float32(blocks+1) < f })
		if index < 0 || index >= len(gpuDeviceBufferTypes) {
Michael Yang's avatar
Michael Yang committed
144
			return cpuDeviceBufferType
145
146
147
		}

		return gpuDeviceBufferTypes[index]
148
149
	}

150
	layers := make([]deviceBufferType, blocks)
151
	for i := range layers {
152
		layers[i] = assignLayer(i)
153
154
	}

155
	output := assignLayer(blocks)
156
157
158
159
160

	maxTensors := len(meta.Tensors().Items())
	maxTensors += 1
	maxTensors += blocks * 2

161
162
163
164
165
166
167
	type tensor struct {
		source *fs.Tensor
		target string
	}

	targets := make(map[string][]string)

168
	ctxs := make(map[*C.struct_ggml_backend_buffer_type]*C.struct_ggml_context)
169
	createTensor := func(t tensor, bts []*C.struct_ggml_backend_buffer_type) *C.struct_ggml_tensor {
170
171
172
173
174
175
176
		for _, bt := range bts {
			if _, ok := ctxs[bt]; !ok {
				ctxs[bt] = C.ggml_init(C.struct_ggml_init_params{
					mem_size: C.ggml_tensor_overhead() * C.size_t(maxTensors),
					no_alloc: true,
				})
			}
Michael Yang's avatar
Michael Yang committed
177

178
179
180
181
182
183
184
185
			targets[t.source.Name] = append(targets[t.source.Name], t.target)

			name := t.source.Name
			if t.target != "" {
				name = t.target
			}

			cname := C.CString(name)
Michael Yang's avatar
Michael Yang committed
186
			defer C.free(unsafe.Pointer(cname))
187
188
189
190
			if tt := C.ggml_get_tensor(ctxs[bt], cname); tt != nil {
				return tt
			}

191
			tt := C.ggml_new_tensor(ctxs[bt], t.source.Kind, C.int(len(t.source.Shape)), (*C.int64_t)(unsafe.Pointer(&t.source.Shape[0])))
Michael Yang's avatar
Michael Yang committed
192
193
			C.ggml_set_name(tt, cname)

194
			slog.Debug("created tensor", "name", name, "shape", t.source.Shape, "dtype", t.source.Kind, "buffer_type", C.GoString(C.ggml_backend_buft_name(bt)))
195
196
197
198
199
			//nolint:staticcheck // TODO: check if buffer type supports this tensor
			return tt
		}

		return nil
Michael Yang's avatar
Michael Yang committed
200
201
	}

202
	contains := func(s string, parts ...string) bool {
203
204
205
206
207
208
209
210
		split := strings.Split(s, ".")
		for _, part := range parts {
			if slices.Contains(split, part) {
				return true
			}
		}

		return false
Michael Yang's avatar
Michael Yang committed
211
212
	}

213
214
	for _, t := range meta.Tensors().Items() {
		switch {
215
		case contains(t.Name, "position_embd", "token_embd", "token_norm_embd", "token_types"):
216
			createTensor(tensor{source: t}, input.bts)
217
		case contains(t.Name, "cls", "output", "output_norm"):
218
			createTensor(tensor{source: t}, output.bts)
219
220
		case strings.HasPrefix(t.Name, "v.") || strings.HasPrefix(t.Name, "mm."):
			createTensor(tensor{source: t}, input.bts)
221
222
223
224
225
226
227
228
229
230
		default:
			if i := func() int {
				if fields := strings.FieldsFunc(t.Name, func(r rune) bool { return !unicode.IsNumber(r) }); len(fields) > 0 {
					if i, err := strconv.Atoi(fields[0]); err == nil {
						return i
					}
				}

				return -1
			}(); i >= 0 {
231
				createTensor(tensor{source: t}, layers[i].bts)
232
			} else {
233
234
235
236
237
				for i, layer := range layers {
					createTensor(tensor{
						source: t,
						target: "blk." + strconv.Itoa(i) + "." + t.Name,
					}, layer.bts)
238
239
240
241
				}
			}
		}
	}
Michael Yang's avatar
Michael Yang committed
242

243
244
245
246
247
248
249
250
251
252
253
254
255
256
	bbs := make(map[*C.struct_ggml_context][]*C.struct_ggml_backend_buffer, len(ctxs))

	for bt, c := range ctxs {
		if C.ggml_get_first_tensor(c) == nil {
			continue
		}

		b := C.ggml_backend_alloc_ctx_tensors_from_buft(c, bt)
		C.ggml_backend_buffer_set_usage(b, C.GGML_BACKEND_BUFFER_USAGE_WEIGHTS)
		bbs[c] = append(bbs[c], b)
	}

	for bs := range maps.Values(bbs) {
		for _, b := range bs {
257
			slog.Info("model weights", "buffer", C.GoString(C.ggml_backend_buffer_name(b)), "size", format.HumanBytes2(uint64(C.ggml_backend_buffer_get_size(b))))
258
259
260
261
262
263
264
265
266
267
268
		}
	}

	tensors := make(map[string]*C.struct_ggml_tensor)
	for _, c := range ctxs {
		for t := C.ggml_get_first_tensor(c); t != nil; t = C.ggml_get_next_tensor(c, t) {
			tensors[C.GoString(C.ggml_get_name(t))] = t
		}
	}

	sr := io.NewSectionReader(r, int64(meta.Tensors().Offset), n-int64(meta.Tensors().Offset))
Michael Yang's avatar
Michael Yang committed
269
	var g errgroup.Group
270
	for _, t := range meta.Tensors().Items() {
271
272
273
274
275
		for _, target := range targets[t.Name] {
			g.Go(func() error {
				if target == "" {
					target = t.Name
				}
276

277
278
279
280
				tt, ok := tensors[target]
				if !ok {
					return fmt.Errorf("unassigned tensor: %s", t.Name)
				}
Michael Yang's avatar
Michael Yang committed
281

282
283
284
285
286
				bts := make([]byte, t.Size())
				n, err := io.ReadFull(io.NewSectionReader(sr, int64(t.Offset), int64(t.Size())), bts)
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
287

288
289
290
				if n != len(bts) {
					return errors.New("short read")
				}
Michael Yang's avatar
Michael Yang committed
291

292
293
294
295
296
297
298
				cname := C.CString(t.Name)
				C.ggml_backend_tensor_set(tt, unsafe.Pointer(&bts[0]), 0, C.size_t(t.Size()))
				C.free(unsafe.Pointer(cname))

				return nil
			})
		}
Michael Yang's avatar
Michael Yang committed
299
300
	}

301
	if g.Wait() != nil {
Michael Yang's avatar
Michael Yang committed
302
303
304
		return nil, err
	}

305
	deviceBackends := make(map[*C.struct_ggml_backend_device]*C.struct_ggml_backend)
306
307
308
309
310
	var backends []*C.struct_ggml_backend
	var bufts []*C.struct_ggml_backend_buffer_type
	for _, d := range append(gpus, append(accels, cpus...)...) {
		b := C.ggml_backend_dev_init(d, nil)
		backends = append(backends, b)
311
		deviceBackends[d] = b
312
313
314
315
316
317
318
319
320
321

		bt := C.ggml_backend_get_default_buffer_type(b)
		if d := C.ggml_backend_get_device(b); C.ggml_backend_dev_type(d) == C.GGML_BACKEND_DEVICE_TYPE_CPU && len(gpus) > 0 {
			if hbt := C.ggml_backend_dev_host_buffer_type(d); hbt != nil {
				bt = hbt
			}
		}

		bufts = append(bufts, bt)

322
		slog.Info("compute graph", "backend", C.GoString(C.ggml_backend_name(b)), "buffer_type", C.GoString(C.ggml_backend_buft_name(bt)))
323
324
325
326

		if C.ggml_backend_is_cpu(b) {
			C.ggml_backend_cpu_set_n_threads(b, C.int(params.NumThreads))
		}
327
328
	}

Michael Yang's avatar
Michael Yang committed
329
	return &Backend{
330
		flashAttention: params.FlashAttention,
331
332
		meta:           meta,
		tensors:        tensors,
333
334
335
336
337
338
339
		sched: C.ggml_backend_sched_new(
			(*C.ggml_backend_t)(unsafe.Pointer(&backends[0])),
			(*C.ggml_backend_buffer_type_t)(unsafe.Pointer(&bufts[0])),
			C.int(len(backends)),
			C.size_t(max(8192, len(meta.Tensors().Items())*5)),
			true,
		),
340
341
342
343
344
345
346
347
348
		input:  deviceBackends[input.d],
		output: deviceBackends[output.d],
		layers: func() map[int]*C.struct_ggml_backend {
			m := make(map[int]*C.struct_ggml_backend)
			for i, layer := range layers {
				m[i] = deviceBackends[layer.d]
			}
			return m
		}(),
Michael Yang's avatar
Michael Yang committed
349
350
351
352
353
354
355
356
357
358
359
360
	}, nil
}

func init() {
	ml.RegisterBackend("ggml", New)
}

func (b *Backend) Config() ml.Config {
	return b.meta.KV()
}

func (b *Backend) Get(name string) ml.Tensor {
361
362
	if t, ok := b.tensors[name]; ok {
		return &Tensor{b: b, t: t}
Michael Yang's avatar
Michael Yang committed
363
364
365
366
367
368
	}

	return nil
}

func (b *Backend) NewContext() ml.Context {
369
370
371
372
	return b.NewContextSize(max(8192, len(b.meta.Tensors().Items())*5))
}

func (b *Backend) NewContextSize(n int) ml.Context {
Michael Yang's avatar
Michael Yang committed
373
	return &Context{
374
		b: b,
375
		ctx: C.ggml_init(C.struct_ggml_init_params{
376
			mem_size: C.size_t(n)*C.ggml_tensor_overhead() + C.ggml_graph_overhead_custom(C.size_t(n), false),
377
378
			no_alloc: true,
		}),
379
		backend:       C.ggml_backend_sched_get_backend(b.sched, 0),
380
381
382
383
		maxGraphNodes: n,
		input:         b.input,
		output:        b.output,
		layers:        b.layers,
Michael Yang's avatar
Michael Yang committed
384
385
386
	}
}

387
func (b *Backend) CacheConfig() ml.CacheConfig {
388
389
390
391
392
	if b.flashAttention {
		return ml.CacheConfig{CachePadding: 256, MaskDType: ml.DTypeF16, MaskBatchPadding: C.GGML_KQ_MASK_PAD}
	} else {
		return ml.CacheConfig{CachePadding: 32, PermutedV: true}
	}
393
394
}

Michael Yang's avatar
Michael Yang committed
395
type Context struct {
396
	b *Backend
Michael Yang's avatar
Michael Yang committed
397

398
	ctx   *C.struct_ggml_context
Michael Yang's avatar
Michael Yang committed
399
	graph *C.struct_ggml_cgraph
400
401

	// backend is the backend used for new tensors
402
	backend *C.struct_ggml_backend
403

404
405
406
407
408
409
410
411
412
	// input is the backend used for inputs
	input *C.struct_ggml_backend

	// output is the backend used for outputs
	output *C.struct_ggml_backend

	// output is the backend used for repeating layers
	layers map[int]*C.struct_ggml_backend

413
	maxGraphNodes int
Michael Yang's avatar
Michael Yang committed
414
415
}

416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
func (c *Context) Input() ml.Context {
	if c.input != nil {
		return &Context{
			b:             c.b,
			ctx:           c.ctx,
			backend:       c.input,
			maxGraphNodes: c.maxGraphNodes,
		}
	}

	return c
}

func (c *Context) Output() ml.Context {
	if c.output != nil {
		return &Context{
			b:             c.b,
			ctx:           c.ctx,
			backend:       c.output,
			maxGraphNodes: c.maxGraphNodes,
		}
	}

	return c
}

func (c *Context) Layer(i int) ml.Context {
	if backend, ok := c.layers[i]; ok {
		return &Context{
			b:             c.b,
			ctx:           c.ctx,
			backend:       backend,
			maxGraphNodes: c.maxGraphNodes,
		}
	}

	return c
}

455
func (c *Context) Forward(tensors ...ml.Tensor) ml.Context {
Michael Yang's avatar
Michael Yang committed
456
	if c.graph == nil {
457
		c.graph = C.ggml_new_graph_custom(c.ctx, C.size_t(c.maxGraphNodes), false)
Michael Yang's avatar
Michael Yang committed
458
459
	}

460
461
462
463
464
	for _, tensor := range tensors {
		C.ggml_build_forward_expand(c.graph, tensor.(*Tensor).t)
	}

	return c
Michael Yang's avatar
Michael Yang committed
465
466
}

467
func (c *Context) Compute(tensors ...ml.Tensor) {
468
	C.ggml_backend_sched_reset(c.b.sched)
469
470
	C.ggml_backend_sched_alloc_graph(c.b.sched, c.graph)
	C.ggml_backend_sched_graph_compute_async(c.b.sched, c.graph)
Michael Yang's avatar
Michael Yang committed
471

472
473
474
	needSync := true
	sync := func() {
		if needSync {
475
			C.ggml_backend_sched_synchronize(c.b.sched)
476
477
478
			needSync = false
		}
	}
Michael Yang's avatar
Michael Yang committed
479

480
481
482
	for _, t := range tensors {
		if C.ggml_nbytes(t.(*Tensor).t) > 0 {
			t.(*Tensor).sync = sync
483
484
		}
	}
Michael Yang's avatar
Michael Yang committed
485
486
}

487
488
func (c *Context) MaxGraphNodes() int {
	return c.maxGraphNodes
Jesse Gross's avatar
Jesse Gross committed
489
490
}

491
492
493
func shapeToGGML(shape []int) *C.int64_t {
	sh := make([]C.int64_t, len(shape))
	for i, s := range shape {
494
		sh[i] = C.int64_t(s)
495
496
497
498
499
	}

	return &sh[0]
}

500
func (c Context) newTensor(dtype ml.DType, shape []int) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
501
502
503
504
505
506
507
508
509
510
511
512
513
	if len(shape) < 1 || len(shape) > 4 {
		panic("unsupported number of dimensions")
	}

	for _, dim := range shape {
		if dim < 1 {
			panic("invalid shape")
		}
	}

	var t *C.struct_ggml_tensor
	switch dtype {
	case ml.DTypeF32:
514
		t = C.ggml_new_tensor(c.ctx, C.GGML_TYPE_F32, C.int(len(shape)), shapeToGGML(shape))
Jesse Gross's avatar
Jesse Gross committed
515
	case ml.DTypeF16:
516
		t = C.ggml_new_tensor(c.ctx, C.GGML_TYPE_F16, C.int(len(shape)), shapeToGGML(shape))
Michael Yang's avatar
Michael Yang committed
517
	case ml.DTypeI32:
518
		t = C.ggml_new_tensor(c.ctx, C.GGML_TYPE_I32, C.int(len(shape)), shapeToGGML(shape))
Michael Yang's avatar
Michael Yang committed
519
520
521
522
	default:
		panic("unsupported dtype")
	}

523
	b := C.ggml_backend_alloc_buffer(c.backend, C.ggml_nbytes(t))
Michael Yang's avatar
Michael Yang committed
524
	C.ggml_backend_tensor_alloc(b, t, C.ggml_backend_buffer_get_base(b))
525
	return &Tensor{b: c.b, t: t}
526
527
528
}

func (c Context) Empty(dtype ml.DType, shape ...int) ml.Tensor {
529
	return c.newTensor(dtype, shape)
530
531
532
}

func (c Context) Zeros(dtype ml.DType, shape ...int) ml.Tensor {
533
	t := c.newTensor(dtype, shape)
534
535
	C.ggml_set_zero(t.(*Tensor).t)
	return t
Michael Yang's avatar
Michael Yang committed
536
537
}

538
func checkShape[S ~[]E, E any](s S, shape ...int) error {
Michael Yang's avatar
Michael Yang committed
539
540
541
542
543
544
	n := len(s)
	for _, v := range shape {
		n /= v
	}

	if n != 1 {
545
		return fmt.Errorf("invalid shape: %v", shape)
Michael Yang's avatar
Michael Yang committed
546
547
	}

548
	return nil
Michael Yang's avatar
Michael Yang committed
549
550
551
}

func (c Context) FromFloatSlice(s []float32, shape ...int) (ml.Tensor, error) {
552
553
554
555
556
557
558
	if err := checkShape(s, shape...); err != nil {
		return nil, err
	}

	t := c.newTensor(ml.DTypeF32, shape)
	C.ggml_backend_tensor_set(t.(*Tensor).t, unsafe.Pointer(&s[0]), 0, C.ggml_nbytes(t.(*Tensor).t))
	return t, nil
Michael Yang's avatar
Michael Yang committed
559
560
561
}

func (c Context) FromIntSlice(s []int32, shape ...int) (ml.Tensor, error) {
562
563
564
565
566
567
568
	if err := checkShape(s, shape...); err != nil {
		return nil, err
	}

	t := c.newTensor(ml.DTypeI32, shape)
	C.ggml_backend_tensor_set(t.(*Tensor).t, unsafe.Pointer(&s[0]), 0, C.ggml_nbytes(t.(*Tensor).t))
	return t, nil
Michael Yang's avatar
Michael Yang committed
569
570
}

571
572
func (c Context) Close() {
	if c.ctx != nil {
573
574
		C.ggml_free(c.ctx)
	}
Michael Yang's avatar
Michael Yang committed
575
576
577
}

type Tensor struct {
578
	b    *Backend
Michael Yang's avatar
Michael Yang committed
579
	t    *C.struct_ggml_tensor
580
	sync func()
Michael Yang's avatar
Michael Yang committed
581
582
583
584
585
586
587
588
589
590
}

func (t *Tensor) LogValue() slog.Value {
	return slog.GroupValue(
		slog.String("name", C.GoString(C.ggml_get_name(t.t))),
		slog.String("type", C.GoString(C.ggml_type_name(t.t._type))),
		slog.Any("shape", t.Shape()),
	)
}

591
592
func (t *Tensor) Dim(n int) int {
	return int(t.t.ne[n])
Michael Yang's avatar
Michael Yang committed
593
594
}

595
596
func (t *Tensor) Stride(n int) int {
	return int(t.t.nb[n])
Michael Yang's avatar
Michael Yang committed
597
598
}

599
600
func (t *Tensor) Shape() []int {
	shape := make([]int, C.ggml_n_dims(t.t))
Michael Yang's avatar
Michael Yang committed
601
602
603
604
605
606
607
	for i := range shape {
		shape[i] = t.Dim(i)
	}

	return shape
}

608
609
610
611
612
613
614
615
616
func (t *Tensor) Bytes() (data []byte) {
	if t.sync != nil {
		data = make([]byte, C.ggml_nbytes(t.t))

		t.sync()
		C.ggml_backend_tensor_get(t.t, unsafe.Pointer(&data[0]), 0, C.ggml_nbytes(t.t))
	}

	return
Michael Yang's avatar
Michael Yang committed
617
618
}

619
620
621
622
623
624
func (t *Tensor) Floats() (data []float32) {
	if t.sync != nil {
		data = make([]float32, C.ggml_nelements(t.t))

		t.sync()
		C.ggml_backend_tensor_get(t.t, unsafe.Pointer(&data[0]), 0, C.ggml_nbytes(t.t))
Michael Yang's avatar
Michael Yang committed
625
626
627
628
629
630
631
632
633
	}

	return
}

func (t *Tensor) DType() ml.DType {
	switch t.t._type {
	case C.GGML_TYPE_F32:
		return ml.DTypeF32
Jesse Gross's avatar
Jesse Gross committed
634
635
	case C.GGML_TYPE_F16:
		return ml.DTypeF16
Michael Yang's avatar
Michael Yang committed
636
637
638
639
640
641
642
643
644
	case C.GGML_TYPE_I32:
		return ml.DTypeI32
	default:
		return ml.DTypeOther
	}
}

func (t *Tensor) Add(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	return &Tensor{
645
		b: t.b,
Michael Yang's avatar
Michael Yang committed
646
647
648
649
650
651
652
653
654
655
656
657
658
659
		t: C.ggml_add(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

func (t *Tensor) Stack(ctx ml.Context, dim int, s ...ml.Tensor) ml.Tensor {
	if len(s) > 0 {
		return t.Concat(ctx, s[0].Stack(ctx, dim, s[1:]...), dim)
	}

	return t
}

func (t *Tensor) Concat(ctx ml.Context, t2 ml.Tensor, dim int) ml.Tensor {
	return &Tensor{
660
		b: t.b,
Michael Yang's avatar
Michael Yang committed
661
662
663
664
665
666
		t: C.ggml_concat(ctx.(*Context).ctx, t.t, t2.(*Tensor).t, C.int(dim)),
	}
}

func (t *Tensor) Contiguous(ctx ml.Context) ml.Tensor {
	return &Tensor{
667
		b: t.b,
Michael Yang's avatar
Michael Yang committed
668
669
670
671
672
673
		t: C.ggml_cont(ctx.(*Context).ctx, t.t),
	}
}

func (t *Tensor) Mul(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	return &Tensor{
674
		b: t.b,
Michael Yang's avatar
Michael Yang committed
675
676
677
678
679
680
		t: C.ggml_mul(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

func (t *Tensor) Mulmat(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	return &Tensor{
681
		b: t.b,
Michael Yang's avatar
Michael Yang committed
682
683
684
685
		t: C.ggml_mul_mat(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

686
687
688
689
690
func (t *Tensor) MulmatFullPrec(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	mul := C.ggml_mul_mat(ctx.(*Context).ctx, t.t, t2.(*Tensor).t)
	C.ggml_mul_mat_set_prec(mul, C.GGML_PREC_F32)

	return &Tensor{
691
		b: t.b,
692
693
694
695
		t: mul,
	}
}

Michael Yang's avatar
Michael Yang committed
696
func (t *Tensor) LayerNorm(ctx ml.Context, w, b ml.Tensor, eps float32) ml.Tensor {
697
	tt := (&Tensor{b: t.b, t: C.ggml_norm(ctx.(*Context).ctx, t.t, C.float(eps))}).Mul(ctx, w)
Michael Yang's avatar
Michael Yang committed
698
699
700
701
702
703
704
705
	if b != nil {
		tt = tt.Add(ctx, b)
	}

	return tt
}

func (t *Tensor) RMSNorm(ctx ml.Context, w ml.Tensor, eps float32) ml.Tensor {
706
	return (&Tensor{b: t.b, t: C.ggml_rms_norm(ctx.(*Context).ctx, t.t, C.float(eps))}).Mul(ctx, w)
Michael Yang's avatar
Michael Yang committed
707
708
}

709
func (t *Tensor) Pad(ctx ml.Context, shape ...int) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
710
711
712
713
714
	if len(shape) != 4 {
		panic("expected 4 dimensions")
	}

	return &Tensor{
715
		b: t.b,
Michael Yang's avatar
Michael Yang committed
716
717
718
719
720
721
722
723
724
725
		t: C.ggml_pad(ctx.(*Context).ctx, t.t, C.int(shape[0]), C.int(shape[1]), C.int(shape[2]), C.int(shape[3])),
	}
}

func (t *Tensor) Permute(ctx ml.Context, shape ...int) ml.Tensor {
	if len(shape) != 4 {
		panic("expected 4 dimensions")
	}

	return &Tensor{
726
		b: t.b,
Michael Yang's avatar
Michael Yang committed
727
728
729
730
731
732
		t: C.ggml_permute(ctx.(*Context).ctx, t.t, C.int(shape[0]), C.int(shape[1]), C.int(shape[2]), C.int(shape[3])),
	}
}

func (t *Tensor) Rows(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	return &Tensor{
733
		b: t.b,
Michael Yang's avatar
Michael Yang committed
734
735
736
737
738
739
		t: C.ggml_get_rows(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

func (t *Tensor) Copy(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	return &Tensor{
740
		b: t.b,
Michael Yang's avatar
Michael Yang committed
741
742
743
744
		t: C.ggml_cpy(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

745
func (t *Tensor) Reshape(ctx ml.Context, shape ...int) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
746
747
748
	switch len(shape) {
	case 1:
		return &Tensor{
749
			b: t.b,
Michael Yang's avatar
Michael Yang committed
750
751
752
753
			t: C.ggml_reshape_1d(ctx.(*Context).ctx, t.t, C.int64_t(shape[0])),
		}
	case 2:
		return &Tensor{
754
			b: t.b,
Michael Yang's avatar
Michael Yang committed
755
756
757
758
			t: C.ggml_reshape_2d(ctx.(*Context).ctx, t.t, C.int64_t(shape[0]), C.int64_t(shape[1])),
		}
	case 3:
		return &Tensor{
759
			b: t.b,
Michael Yang's avatar
Michael Yang committed
760
761
762
763
			t: C.ggml_reshape_3d(ctx.(*Context).ctx, t.t, C.int64_t(shape[0]), C.int64_t(shape[1]), C.int64_t(shape[2])),
		}
	case 4:
		return &Tensor{
764
			b: t.b,
Michael Yang's avatar
Michael Yang committed
765
766
767
768
769
770
771
772
773
			t: C.ggml_reshape_4d(ctx.(*Context).ctx, t.t, C.int64_t(shape[0]), C.int64_t(shape[1]), C.int64_t(shape[2]), C.int64_t(shape[3])),
		}
	default:
		panic("unsupported number of dimensions")
	}
}

func (t *Tensor) Scale(ctx ml.Context, s float64) ml.Tensor {
	return &Tensor{
774
		b: t.b,
Michael Yang's avatar
Michael Yang committed
775
776
777
778
779
780
		t: C.ggml_scale(ctx.(*Context).ctx, t.t, (C.float)(s)),
	}
}

func (t *Tensor) Softmax(ctx ml.Context) ml.Tensor {
	return &Tensor{
781
		b: t.b,
Michael Yang's avatar
Michael Yang committed
782
783
784
785
786
787
		t: C.ggml_soft_max(ctx.(*Context).ctx, t.t),
	}
}

func (t *Tensor) Tanh(ctx ml.Context) ml.Tensor {
	return &Tensor{
788
		b: t.b,
Michael Yang's avatar
Michael Yang committed
789
790
791
792
		t: C.ggml_tanh_inplace(ctx.(*Context).ctx, t.t),
	}
}

793
func (t *Tensor) Unpad(ctx ml.Context, shape ...int) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
794
795
796
797
798
	if len(shape) != 4 {
		panic("expected 4 dimensions")
	}

	return &Tensor{
799
		b: t.b,
Michael Yang's avatar
Michael Yang committed
800
801
802
803
804
805
806
807
		t: C.ggml_unpad(ctx.(*Context).ctx, t.t, C.int(shape[0]), C.int(shape[1]), C.int(shape[2]), C.int(shape[3])),
	}
}

func (t *Tensor) View(ctx ml.Context, offset int, shape ...int) ml.Tensor {
	switch len(shape) {
	case 1:
		return &Tensor{
808
			b: t.b,
Michael Yang's avatar
Michael Yang committed
809
810
811
812
			t: C.ggml_view_1d(ctx.(*Context).ctx, t.t, C.int64_t(shape[0]), C.size_t(offset)),
		}
	case 3:
		return &Tensor{
813
			b: t.b,
Michael Yang's avatar
Michael Yang committed
814
815
816
817
818
819
820
			t: C.ggml_view_2d(ctx.(*Context).ctx, t.t,
				C.int64_t(shape[0]), C.int64_t(shape[2]),
				C.size_t(shape[1]),
				C.size_t(offset)),
		}
	case 5:
		return &Tensor{
821
			b: t.b,
Michael Yang's avatar
Michael Yang committed
822
823
824
825
826
827
828
			t: C.ggml_view_3d(ctx.(*Context).ctx, t.t,
				C.int64_t(shape[0]), C.int64_t(shape[2]), C.int64_t(shape[4]),
				C.size_t(shape[1]), C.size_t(shape[3]),
				C.size_t(offset)),
		}
	case 7:
		return &Tensor{
829
			b: t.b,
Michael Yang's avatar
Michael Yang committed
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
			t: C.ggml_view_4d(ctx.(*Context).ctx, t.t,
				C.int64_t(shape[0]), C.int64_t(shape[2]), C.int64_t(shape[4]), C.int64_t(shape[6]),
				C.size_t(shape[1]), C.size_t(shape[3]), C.size_t(shape[5]),
				C.size_t(offset)),
		}
	default:
		panic("unsupported number of dimensions")
	}
}

const (
	ropeTypeNorm C.int = iota
)

func (t *Tensor) RoPE(ctx ml.Context, positionIDs, ropeFactors ml.Tensor, ropeDim uint32, ropeBase, ropeScale float32) ml.Tensor {
	if ropeFactors == nil {
846
		ropeFactors = &Tensor{b: t.b}
Michael Yang's avatar
Michael Yang committed
847
848
	}

Jesse Gross's avatar
Jesse Gross committed
849
850
851
852
853
	dequant := t.t
	if C.ggml_is_quantized(t.t._type) {
		dequant = C.ggml_cast(ctx.(*Context).ctx, t.t, C.GGML_TYPE_F32)
	}

Michael Yang's avatar
Michael Yang committed
854
	return &Tensor{
855
		b: t.b,
Michael Yang's avatar
Michael Yang committed
856
		t: C.ggml_rope_ext(
Jesse Gross's avatar
Jesse Gross committed
857
			ctx.(*Context).ctx, dequant, positionIDs.(*Tensor).t, ropeFactors.(*Tensor).t,
Michael Yang's avatar
Michael Yang committed
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
			C.int(ropeDim),
			131072,       // YaRN n_ctx_train
			ropeTypeNorm, // ROPE_TYPE_NORM
			C.float(ropeBase),
			C.float(ropeScale),
			0.,  // YaRN ext_factor
			1.,  // YaRN attn_factor
			32., // YaRN beta_fast
			1.,  // YaRN beta_slow
		),
	}
}

func (t *Tensor) GELU(ctx ml.Context) ml.Tensor {
	return &Tensor{
873
		b: t.b,
Michael Yang's avatar
Michael Yang committed
874
875
876
877
878
879
		t: C.ggml_gelu_inplace(ctx.(*Context).ctx, t.t),
	}
}

func (t *Tensor) SILU(ctx ml.Context) ml.Tensor {
	return &Tensor{
880
		b: t.b,
Michael Yang's avatar
Michael Yang committed
881
882
883
884
885
886
		t: C.ggml_silu_inplace(ctx.(*Context).ctx, t.t),
	}
}

func (t *Tensor) Conv2D(ctx ml.Context, t2 ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor {
	return &Tensor{
887
		b: t.b,
Michael Yang's avatar
Michael Yang committed
888
889
890
		t: C.ggml_conv_2d(ctx.(*Context).ctx, t.t, t2.(*Tensor).t, C.int(s0), C.int(s1), C.int(p0), C.int(p1), C.int(d0), C.int(d1)),
	}
}
891

892
893
894
895
896
897
func (t *Tensor) ScaledDotProductAttention(ctx ml.Context, key, value, mask ml.Tensor, scale float64) ml.Tensor {
	var kqMask *C.struct_ggml_tensor
	if mask != nil {
		kqMask = mask.(*Tensor).t
	}

898
899
900
	query := t.Permute(ctx, 0, 2, 1, 3)
	key = key.Permute(ctx, 0, 2, 1, 3)

901
902
	if t.b.flashAttention {
		value = value.Permute(ctx, 0, 2, 1, 3)
903

904
905
906
907
908
909
910
911
912
913
914
915
916
		kqv := C.ggml_flash_attn_ext(ctx.(*Context).ctx, query.(*Tensor).t, key.(*Tensor).t, value.(*Tensor).t, kqMask, C.float(scale), 0, 0)
		C.ggml_flash_attn_ext_set_prec(kqv, C.GGML_PREC_F32)
		return &Tensor{b: t.b, t: kqv}
	} else {
		kq := key.MulmatFullPrec(ctx, query)
		kq = &Tensor{
			b: t.b,
			t: C.ggml_soft_max_ext(ctx.(*Context).ctx, kq.(*Tensor).t, kqMask, C.float(scale), 0),
		}

		kqv := value.Mulmat(ctx, kq)
		return kqv.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
	}
917
}