ggml.go 27.9 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
	"context"
13
	"errors"
Michael Yang's avatar
Michael Yang committed
14
15
16
	"fmt"
	"io"
	"log/slog"
17
	"maps"
Michael Yang's avatar
Michael Yang committed
18
	"os"
19
	"runtime"
20
21
22
	"slices"
	"strconv"
	"strings"
23
	"sync/atomic"
24
	"unicode"
Michael Yang's avatar
Michael Yang committed
25
26
27
	"unsafe"

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

Michael Yang's avatar
Michael Yang committed
35
36
37
38
39
func devices() []*C.struct_ggml_backend_device {
	ggml.OnceLoad()
	ds := make([]*C.struct_ggml_backend_device, C.ggml_backend_dev_count())
	for i := range ds {
		ds[i] = C.ggml_backend_dev_get(C.size_t(i))
Michael Yang's avatar
Michael Yang committed
40
	}
Michael Yang's avatar
Michael Yang committed
41
42

	return ds
43
}
Michael Yang's avatar
Michael Yang committed
44
45

type Backend struct {
46
47
48
49
50
51
	meta *fsggml.GGML

	sched         *C.struct_ggml_backend_sched
	schedBackends []*C.struct_ggml_backend
	schedBufts    []*C.struct_ggml_backend_buffer_type

52
	tensors map[string]*C.struct_ggml_tensor
Michael Yang's avatar
Michael Yang committed
53
54

	// input is the backend used for inputs
55
	input *C.struct_ggml_backend_buffer_type
Michael Yang's avatar
Michael Yang committed
56
57

	// layers is the backend used for repeating layers
58
	layers map[int]*C.struct_ggml_backend_buffer_type
59

60
	flashAttention bool
Michael Yang's avatar
Michael Yang committed
61
62
63

	// maxGraphNodes is the maximum allowed number of graph nodes in this scheduler
	maxGraphNodes int
Michael Yang's avatar
Michael Yang committed
64
65
}

66
func New(ctx context.Context, r *os.File, params ml.BackendParams) (ml.Backend, error) {
67
	meta, n, err := fsggml.Decode(r, -1)
Michael Yang's avatar
Michael Yang committed
68
69
70
71
72
73
74
75
76
77
78
79
80
81
	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()),
	)

82
	type deviceBufferType struct {
83
84
85
86
87
		d   *C.struct_ggml_backend_device
		bts []*C.struct_ggml_backend_buffer_type
	}

	var cpus, accels, gpus []*C.struct_ggml_backend_device
Michael Yang's avatar
Michael Yang committed
88
	for _, d := range devices() {
89
90
		switch C.ggml_backend_dev_type(d) {
		case C.GGML_BACKEND_DEVICE_TYPE_CPU:
91
92
93
94
			if len(cpus) == 0 {
				// only the first cpu device should be used
				cpus = append(cpus, d)
			}
95
96
		case C.GGML_BACKEND_DEVICE_TYPE_ACCEL:
			accels = append(accels, d)
Michael Yang's avatar
Michael Yang committed
97
		case C.GGML_BACKEND_DEVICE_TYPE_GPU:
98
			gpus = append(gpus, d)
Michael Yang's avatar
Michael Yang committed
99
100
101
		}
	}

Michael Yang's avatar
Michael Yang committed
102
	// create list of buffer types for the cpu
Michael Yang's avatar
Michael Yang committed
103
	cpuDeviceBufferType := deviceBufferType{d: C.ggml_backend_dev_by_type(C.GGML_BACKEND_DEVICE_TYPE_CPU)}
104
105
106
107
	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
108
			cpuDeviceBufferType.bts = append(cpuDeviceBufferType.bts, C.ggml_backend_dev_buffer_type(d))
Michael Yang's avatar
Michael Yang committed
109
		}
110
111
	}

Michael Yang's avatar
Michael Yang committed
112
	// create list of buffer types for each gpu
113
	var gpuDeviceBufferTypes []deviceBufferType
114
115
	for _, d := range gpus {
		bt := C.ggml_backend_dev_buffer_type(d)
116
		gpuDeviceBufferTypes = append(gpuDeviceBufferTypes, deviceBufferType{
117
			d:   d,
Michael Yang's avatar
Michael Yang committed
118
			bts: append([]*C.struct_ggml_backend_buffer_type{bt}, cpuDeviceBufferType.bts...),
119
		})
Michael Yang's avatar
Michael Yang committed
120
121
	}

Michael Yang's avatar
Michael Yang committed
122
123
124
125
126
	useDefaultSplit := true
	for _, s := range params.TensorSplit {
		if s != 0 {
			useDefaultSplit = false
			break
127
		}
Michael Yang's avatar
Michael Yang committed
128
	}
129

Michael Yang's avatar
Michael Yang committed
130
131
132
133
	// calculate splits
	splits := make([]float32, len(gpus))
	if useDefaultSplit {
		// default: split on free memory
134
135
136
137
138
		for i := range splits {
			var free, total C.size_t
			C.ggml_backend_dev_memory(gpus[i], &free, &total)
			splits[i] = float32(free)
		}
Michael Yang's avatar
Michael Yang committed
139
140
	} else {
		splits = params.TensorSplit
141
142
143
	}

	var sum float32
Michael Yang's avatar
Michael Yang committed
144
	// cumulative sum of all splits
145
146
147
148
149
	for i := range splits {
		sum += splits[i]
		splits[i] = sum
	}

Michael Yang's avatar
Michael Yang committed
150
	// normalize splits
151
	for i := range splits {
152
		splits[i] /= sum
153
154
	}

Michael Yang's avatar
Michael Yang committed
155
	// inputs always use cpu
Michael Yang's avatar
Michael Yang committed
156
	input := cpuDeviceBufferType
157

158
	blocks := int(meta.KV().BlockCount())
Michael Yang's avatar
Michael Yang committed
159
160
161
162

	// define a range of gpu layers. anything outside of this range is assigned to the cpu
	gpuRangeStart := max(0, blocks-params.NumGPULayers)
	gpuRangeStop := min(gpuRangeStart+params.NumGPULayers, blocks+1)
Michael Yang's avatar
Michael Yang committed
163
	assignLayer := func(i int) deviceBufferType {
Michael Yang's avatar
Michael Yang committed
164
		if i < gpuRangeStart || i >= gpuRangeStop {
Michael Yang's avatar
Michael Yang committed
165
			return cpuDeviceBufferType
166
		}
167

Michael Yang's avatar
Michael Yang committed
168
		index := slices.IndexFunc(splits, func(f float32) bool { return float32(i-gpuRangeStart)/float32(gpuRangeStop-gpuRangeStart) < f })
169
		if index < 0 || index >= len(gpuDeviceBufferTypes) {
Michael Yang's avatar
Michael Yang committed
170
			return cpuDeviceBufferType
171
172
173
		}

		return gpuDeviceBufferTypes[index]
174
175
	}

Michael Yang's avatar
Michael Yang committed
176
	// repeating layers are assigned based on their index in reverse order, e.g. i / (block_count + 1)
177
	layers := make([]deviceBufferType, blocks)
178
	for i := range layers {
179
		layers[i] = assignLayer(i)
180
181
	}

Michael Yang's avatar
Michael Yang committed
182
	// outputs are assigned iff allowed by splits and configured number of gpu layers
183
	output := assignLayer(blocks)
184
185
186

	maxTensors := len(meta.Tensors().Items())
	maxTensors += 1
Michael Yang's avatar
Michael Yang committed
187
	// each layer has at most 2 extra tensors for rope operations
188
189
	maxTensors += blocks * 2

190
	type tensor struct {
191
		source *fsggml.Tensor
192
193
194
		target string
	}

Michael Yang's avatar
Michael Yang committed
195
	// some tensors are mapped to different names so keep a list
196
197
	targets := make(map[string][]string)

Michael Yang's avatar
Michael Yang committed
198
	// contexts are shared by tensors of the same buffer type
199
	ctxs := make(map[*C.struct_ggml_backend_buffer_type]*C.struct_ggml_context)
200
	createTensor := func(t tensor, bts []*C.struct_ggml_backend_buffer_type) *C.struct_ggml_tensor {
201
202
203
204
205
206
207
		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
208

209
210
211
212
213
214
215
216
			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
217
			defer C.free(unsafe.Pointer(cname))
218
219
220
221
			if tt := C.ggml_get_tensor(ctxs[bt], cname); tt != nil {
				return tt
			}

222
			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
223
224
			C.ggml_set_name(tt, cname)

225
			slog.Debug("created tensor", "name", name, "shape", t.source.Shape, "dtype", t.source.Kind, "buffer_type", C.GoString(C.ggml_backend_buft_name(bt)))
226
227
228
229
230
			//nolint:staticcheck // TODO: check if buffer type supports this tensor
			return tt
		}

		return nil
Michael Yang's avatar
Michael Yang committed
231
232
	}

233
	contains := func(s string, parts ...string) bool {
234
235
236
237
238
239
240
241
		split := strings.Split(s, ".")
		for _, part := range parts {
			if slices.Contains(split, part) {
				return true
			}
		}

		return false
Michael Yang's avatar
Michael Yang committed
242
243
	}

244
245
	for _, t := range meta.Tensors().Items() {
		switch {
246
		case contains(t.Name, "position_embd", "token_embd", "token_norm_embd", "token_types"):
247
			createTensor(tensor{source: t}, input.bts)
Michael Yang's avatar
Michael Yang committed
248
249
250
			if _, ok := meta.Tensors().GroupLayers()["output"]; !ok && t.Name == "token_embd.weight" {
				createTensor(tensor{source: t, target: "output.weight"}, output.bts)
			}
251
		case contains(t.Name, "cls", "output", "output_norm"):
252
			createTensor(tensor{source: t}, output.bts)
253
		case strings.HasPrefix(t.Name, "v.") || strings.HasPrefix(t.Name, "mm."):
Michael Yang's avatar
Michael Yang committed
254
			// TODO: assign vision tensors to the gpu if possible
Michael Yang's avatar
Michael Yang committed
255
			createTensor(tensor{source: t}, output.bts)
Michael Yang's avatar
Michael Yang committed
256
257
258
259
260
261
262
263
		case contains(t.Name, "rope_freqs", "rope_factors_long", "rope_factors_short"):
			// these tensors should be repeated per layer
			for i, layer := range layers {
				createTensor(tensor{
					source: t,
					target: "blk." + strconv.Itoa(i) + "." + t.Name,
				}, layer.bts)
			}
264
		default:
Michael Yang's avatar
Michael Yang committed
265
266
267
268
			layerIndex := -1
			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 {
					layerIndex = i
269
				}
Michael Yang's avatar
Michael Yang committed
270
			}
271

Michael Yang's avatar
Michael Yang committed
272
273
			if layerIndex >= 0 {
				createTensor(tensor{source: t}, layers[layerIndex].bts)
274
			} else {
Michael Yang's avatar
Michael Yang committed
275
276
				// load all other tensors on the cpu
				createTensor(tensor{source: t}, input.bts)
277
278
279
			}
		}
	}
Michael Yang's avatar
Michael Yang committed
280

Michael Yang's avatar
Michael Yang committed
281
282
	// allocate buffers for each context
	bbs := make(map[*C.struct_ggml_context]*C.struct_ggml_backend_buffer, len(ctxs))
283
284
285
286
287
288
	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)
289
290
291
292
		if b == nil {
			return nil, fmt.Errorf("unable to allocate memory from device %v for model weights", C.GoString(C.ggml_backend_buft_name(bt)))
		}

293
		C.ggml_backend_buffer_set_usage(b, C.GGML_BACKEND_BUFFER_USAGE_WEIGHTS)
Michael Yang's avatar
Michael Yang committed
294
		bbs[c] = b
295
296
297
	}

	for bs := range maps.Values(bbs) {
Michael Yang's avatar
Michael Yang committed
298
		slog.Info("model weights", "buffer", C.GoString(C.ggml_backend_buffer_name(bs)), "size", format.HumanBytes2(uint64(C.ggml_backend_buffer_get_size(bs))))
299
300
	}

Michael Yang's avatar
Michael Yang committed
301
	// map tensor names to tensors for easy lookup later
302
303
304
305
306
307
308
	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
		}
	}

309
310
311
312
313
	var doneBytes atomic.Uint64
	totalBytes := uint64(n) - meta.Tensors().Offset

	g, ctx := errgroup.WithContext(ctx)
	g.SetLimit(runtime.GOMAXPROCS(0))
314
	for _, t := range meta.Tensors().Items() {
315
316
317
318
		g.Go(func() error {
			tts := make([]*C.struct_ggml_tensor, max(1, len(targets[t.Name])))
			for i := range tts {
				target := targets[t.Name][i]
319
320
321
				if target == "" {
					target = t.Name
				}
322

323
324
325
326
				tt, ok := tensors[target]
				if !ok {
					return fmt.Errorf("unassigned tensor: %s", t.Name)
				}
Michael Yang's avatar
Michael Yang committed
327

328
329
330
				tts[i] = tt
			}

331
332
333
334
			// Create a new FD for each goroutine so that each FD is read sequentially, rather than
			// seeking around within an FD shared between all goroutines.
			file, err := os.Open(r.Name())
			if err != nil {
Jesse Gross's avatar
Jesse Gross committed
335
				slog.Warn("file open error", "file", r.Name(), "error", err)
336
337
338
339
				return err
			}
			defer file.Close()
			sr := io.NewSectionReader(file, int64(meta.Tensors().Offset+t.Offset), int64(t.Size()))
340
341
342
343
344
345
			bts := make([]byte, 128*format.KibiByte)

			var s uint64
			for s < t.Size() {
				n, err := io.ReadFull(sr, bts[:min(len(bts), int(t.Size()-s))])
				if err != nil {
Jesse Gross's avatar
Jesse Gross committed
346
					slog.Warn("file read error", "file", r.Name(), "error", err)
347
					return err
348
				}
Michael Yang's avatar
Michael Yang committed
349

350
351
				for _, tt := range tts {
					C.ggml_backend_tensor_set(tt, unsafe.Pointer(&bts[0]), C.size_t(s), C.size_t(n))
352
				}
Michael Yang's avatar
Michael Yang committed
353

354
355
356
357
358
359
360
361
362
363
				s += uint64(n)

				if params.Progress != nil {
					done := doneBytes.Add(uint64(n))
					params.Progress(float32(done) / float32(totalBytes))
				}
			}

			return nil
		})
Michael Yang's avatar
Michael Yang committed
364
365
	}

366
367
368
369
370
371
372
373
	// start a goroutine to cancel the errgroup if the parent context is done
	go func() {
		<-ctx.Done()
		g.Go(func() error {
			return ctx.Err()
		})
	}()

374
	if err := g.Wait(); err != nil {
Michael Yang's avatar
Michael Yang committed
375
376
377
		return nil, err
	}

378
379
	// map devices to backend buffer types so new tensors can be assigned to the correct device
	deviceBufferTypes := make(map[*C.struct_ggml_backend_device]*C.struct_ggml_backend_buffer_type)
Michael Yang's avatar
Michael Yang committed
380
381
382
383

	// create backends and buffer types used for the compute graph scheduler
	var schedBackends []*C.struct_ggml_backend
	var schedBufts []*C.struct_ggml_backend_buffer_type
384
385
386
387
	for _, d := range append(gpus, append(accels, cpus...)...) {
		b := C.ggml_backend_dev_init(d, nil)
		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 {
388
389
			// use the first gpu host buffer type for gpu if possible
			if hbt := C.ggml_backend_dev_host_buffer_type(gpus[0]); hbt != nil {
390
391
392
393
				bt = hbt
			}
		}

394
395
396
		deviceBufferTypes[d] = bt

		schedBackends = append(schedBackends, b)
Michael Yang's avatar
Michael Yang committed
397
		schedBufts = append(schedBufts, bt)
398

399
		if C.ggml_backend_is_cpu(b) {
Michael Yang's avatar
Michael Yang committed
400
			// set number of threads for cpu backend
Michael Yang's avatar
Michael Yang committed
401
			C.ggml_backend_cpu_set_n_threads(b, C.int(Threads(params.NumThreads)))
402
		}
403
404
	}

Michael Yang's avatar
Michael Yang committed
405
	maxGraphNodes := max(8192, len(meta.Tensors().Items())*5)
Michael Yang's avatar
Michael Yang committed
406
	return &Backend{
407
		flashAttention: params.FlashAttention,
408
409
		meta:           meta,
		tensors:        tensors,
410
		sched: C.ggml_backend_sched_new(
Michael Yang's avatar
Michael Yang committed
411
412
413
414
			(*C.ggml_backend_t)(unsafe.Pointer(&schedBackends[0])),
			(*C.ggml_backend_buffer_type_t)(unsafe.Pointer(&schedBufts[0])),
			C.int(len(schedBackends)),
			C.size_t(maxGraphNodes),
415
			C._Bool(len(gpus) > 1 && slices.Contains(gpus, output.d)),
416
		),
417
418
419
		schedBackends: schedBackends,
		schedBufts:    schedBufts,
		input:         deviceBufferTypes[input.d],
420
421
		layers: func() map[int]*C.struct_ggml_backend_buffer_type {
			m := make(map[int]*C.struct_ggml_backend_buffer_type)
422
			for i, layer := range layers {
423
				m[i] = deviceBufferTypes[layer.d]
424
425
426
			}
			return m
		}(),
Michael Yang's avatar
Michael Yang committed
427
		maxGraphNodes: maxGraphNodes,
Michael Yang's avatar
Michael Yang committed
428
429
430
431
432
433
434
	}, nil
}

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

435
func (b *Backend) Config() fs.Config {
Michael Yang's avatar
Michael Yang committed
436
437
438
439
	return b.meta.KV()
}

func (b *Backend) Get(name string) ml.Tensor {
440
441
	if t, ok := b.tensors[name]; ok {
		return &Tensor{b: b, t: t}
Michael Yang's avatar
Michael Yang committed
442
443
444
445
446
447
	}

	return nil
}

func (b *Backend) NewContext() ml.Context {
Michael Yang's avatar
Michael Yang committed
448
	return b.NewContextSize(b.maxGraphNodes)
449
450
451
}

func (b *Backend) NewContextSize(n int) ml.Context {
Jesse Gross's avatar
Jesse Gross committed
452
453
454
455
	if n > b.maxGraphNodes {
		panic(fmt.Errorf("requested number of graph nodes (%v) for new context exceeds maximum (%v)", n, b.maxGraphNodes))
	}

Michael Yang's avatar
Michael Yang committed
456
	return &Context{
457
458
		b:             b,
		maxGraphNodes: n,
459
		ctx: C.ggml_init(C.struct_ggml_init_params{
460
			mem_size: C.size_t(n)*C.ggml_tensor_overhead() + C.ggml_graph_overhead_custom(C.size_t(n), false),
461
462
			no_alloc: true,
		}),
Michael Yang's avatar
Michael Yang committed
463
464
465
	}
}

466
func (b *Backend) CacheConfig() ml.CacheConfig {
467
468
469
470
471
	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}
	}
472
473
}

Michael Yang's avatar
Michael Yang committed
474
type Context struct {
475
	b *Backend
Michael Yang's avatar
Michael Yang committed
476

477
	ctx   *C.struct_ggml_context
Michael Yang's avatar
Michael Yang committed
478
	graph *C.struct_ggml_cgraph
479

480
481
	// buft is the buffer type used for new tensors
	buft *C.struct_ggml_backend_buffer_type
482

Michael Yang's avatar
Michael Yang committed
483
	// maxGraphNodes is the maximum allowed number of graph nodes in this context
484
	maxGraphNodes int
Michael Yang's avatar
Michael Yang committed
485
486
}

487
func (c *Context) Input() ml.Context {
Michael Yang's avatar
Michael Yang committed
488
	if c.b.input != nil {
489
490
491
		return &Context{
			b:             c.b,
			ctx:           c.ctx,
492
			buft:          c.b.input,
493
494
495
496
			maxGraphNodes: c.maxGraphNodes,
		}
	}

497
	return c
498
499
}

500
func (c *Context) Layer(i int) ml.Context {
501
	if buft, ok := c.b.layers[i]; ok {
502
503
504
		return &Context{
			b:             c.b,
			ctx:           c.ctx,
505
			buft:          buft,
506
507
508
509
			maxGraphNodes: c.maxGraphNodes,
		}
	}

510
	return c
511
512
}

513
func (c *Context) Forward(tensors ...ml.Tensor) ml.Context {
Michael Yang's avatar
Michael Yang committed
514
	if c.graph == nil {
515
		c.graph = C.ggml_new_graph_custom(c.ctx, C.size_t(c.maxGraphNodes), false)
Michael Yang's avatar
Michael Yang committed
516
517
	}

518
519
520
521
522
	for _, tensor := range tensors {
		C.ggml_build_forward_expand(c.graph, tensor.(*Tensor).t)
	}

	return c
Michael Yang's avatar
Michael Yang committed
523
524
}

525
func (c *Context) Compute(tensors ...ml.Tensor) {
526
	C.ggml_backend_sched_graph_compute_async(c.b.sched, c.graph)
Michael Yang's avatar
Michael Yang committed
527
	C.ggml_backend_sched_reset(c.b.sched)
Michael Yang's avatar
Michael Yang committed
528

529
530
531
	needSync := true
	sync := func() {
		if needSync {
532
			C.ggml_backend_sched_synchronize(c.b.sched)
533
534
535
			needSync = false
		}
	}
Michael Yang's avatar
Michael Yang committed
536

537
538
539
	for _, t := range tensors {
		if C.ggml_nbytes(t.(*Tensor).t) > 0 {
			t.(*Tensor).sync = sync
540
541
		}
	}
Michael Yang's avatar
Michael Yang committed
542
543
}

544
func (c *Context) Reserve() error {
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
	if !C.ggml_backend_sched_reserve(c.b.sched, c.graph) {
		C.ggml_backend_sched_reset(c.b.sched)
		return errors.New("failed to reserve graph")
	}

	slog.Debug("compute graph", "nodes", C.ggml_graph_n_nodes(c.graph), "splits", C.ggml_backend_sched_get_n_splits(c.b.sched))
	for i := range c.b.schedBackends {
		size := C.ggml_backend_sched_get_buffer_size(c.b.sched, c.b.schedBackends[i])
		slog.Info("compute graph", "backend", C.GoString(C.ggml_backend_name(c.b.schedBackends[i])), "buffer_type", C.GoString(C.ggml_backend_buft_name(c.b.schedBufts[i])),
			"size", format.HumanBytes2(uint64(size)))
	}

	C.ggml_backend_sched_reset(c.b.sched)

	return nil
}

562
func (c *Context) MaxGraphNodes() int {
563
	return c.maxGraphNodes
Jesse Gross's avatar
Jesse Gross committed
564
565
}

566
567
568
func shapeToGGML(shape []int) *C.int64_t {
	sh := make([]C.int64_t, len(shape))
	for i, s := range shape {
569
		sh[i] = C.int64_t(s)
570
571
572
573
574
	}

	return &sh[0]
}

575
576
577
578
func pad(length, pad C.size_t) C.size_t {
	return ((length + pad - 1) / pad) * pad
}

579
func (c *Context) newTensor(dtype ml.DType, shape []int) (ml.Tensor, error) {
580
	if c.buft == nil {
581
		panic("set Input or Layer before creating tensors")
582
583
	}

Michael Yang's avatar
Michael Yang committed
584
585
586
587
588
589
	var cdtype uint32
	switch dtype {
	case ml.DTypeF32:
		cdtype = C.GGML_TYPE_F32
	case ml.DTypeF16:
		cdtype = C.GGML_TYPE_F16
590
591
592
593
	case ml.DTypeQ80:
		cdtype = C.GGML_TYPE_Q8_0
	case ml.DTypeQ40:
		cdtype = C.GGML_TYPE_Q4_0
Michael Yang's avatar
Michael Yang committed
594
595
596
597
598
599
	case ml.DTypeI32:
		cdtype = C.GGML_TYPE_I32
	default:
		panic("unsupported dtype")
	}

Jesse Gross's avatar
Jesse Gross committed
600
	if len(shape) < 1 || shape[0] == 0 {
Michael Yang's avatar
Michael Yang committed
601
		var shape C.int64_t = 0
602
		return &Tensor{b: c.b, t: C.ggml_new_tensor(c.ctx, cdtype, 1, &shape)}, nil
Michael Yang's avatar
Michael Yang committed
603
	} else if len(shape) > 4 {
Michael Yang's avatar
Michael Yang committed
604
605
606
607
608
609
610
611
612
		panic("unsupported number of dimensions")
	}

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

Michael Yang's avatar
Michael Yang committed
613
	t := C.ggml_new_tensor(c.ctx, cdtype, C.int(len(shape)), shapeToGGML(shape))
614
615
	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)
616
617
618
619
	if b == nil {
		return nil, fmt.Errorf("unable to allocate %v from device %v for new tensor", format.HumanBytes2(uint64(size)), C.GoString(C.ggml_backend_buft_name(c.buft)))
	}

Michael Yang's avatar
Michael Yang committed
620
	C.ggml_backend_tensor_alloc(b, t, C.ggml_backend_buffer_get_base(b))
621
	return &Tensor{b: c.b, t: t}, nil
622
623
}

624
func (c *Context) Empty(dtype ml.DType, shape ...int) ml.Tensor {
625
626
627
628
629
630
	t, err := c.newTensor(dtype, shape)
	if err != nil {
		panic(err)
	}

	return t
631
632
}

633
func (c *Context) Zeros(dtype ml.DType, shape ...int) ml.Tensor {
634
635
636
637
638
	t, err := c.newTensor(dtype, shape)
	if err != nil {
		panic(err)
	}

639
640
	C.ggml_set_zero(t.(*Tensor).t)
	return t
Michael Yang's avatar
Michael Yang committed
641
642
}

643
func checkShape[S ~[]E, E any](s S, shape ...int) error {
Michael Yang's avatar
Michael Yang committed
644
	n := len(s)
Jesse Gross's avatar
Jesse Gross committed
645
646
647
648
649

	if n == 0 {
		return nil
	}

Michael Yang's avatar
Michael Yang committed
650
651
652
653
654
	for _, v := range shape {
		n /= v
	}

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

658
	return nil
Michael Yang's avatar
Michael Yang committed
659
660
}

661
func (c *Context) FromFloatSlice(s []float32, shape ...int) (ml.Tensor, error) {
Jesse Gross's avatar
Jesse Gross committed
662
	if err := checkShape(s, shape...); err != nil {
663
664
665
		return nil, err
	}

666
667
668
669
670
	t, err := c.newTensor(ml.DTypeF32, shape)
	if err != nil {
		return nil, err
	}

Jesse Gross's avatar
Jesse Gross committed
671
672
673
674
	if len(s) > 0 {
		C.ggml_backend_tensor_set(t.(*Tensor).t, unsafe.Pointer(&s[0]), 0, C.ggml_nbytes(t.(*Tensor).t))
	}

675
	return t, nil
Michael Yang's avatar
Michael Yang committed
676
677
}

678
func (c *Context) FromIntSlice(s []int32, shape ...int) (ml.Tensor, error) {
Jesse Gross's avatar
Jesse Gross committed
679
	if err := checkShape(s, shape...); err != nil {
680
681
682
		return nil, err
	}

683
684
685
686
687
	t, err := c.newTensor(ml.DTypeI32, shape)
	if err != nil {
		return nil, err
	}

Jesse Gross's avatar
Jesse Gross committed
688
689
690
691
	if len(s) > 0 {
		C.ggml_backend_tensor_set(t.(*Tensor).t, unsafe.Pointer(&s[0]), 0, C.ggml_nbytes(t.(*Tensor).t))
	}

692
	return t, nil
Michael Yang's avatar
Michael Yang committed
693
694
}

Michael Yang's avatar
Michael Yang committed
695
696
func (c *Context) Close() {
	if c != nil {
697
698
		C.ggml_free(c.ctx)
	}
Michael Yang's avatar
Michael Yang committed
699
700
701
}

type Tensor struct {
702
	b    *Backend
Michael Yang's avatar
Michael Yang committed
703
	t    *C.struct_ggml_tensor
704
	sync func()
Michael Yang's avatar
Michael Yang committed
705
706
707
708
709
710
711
712
713
714
}

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()),
	)
}

715
716
func (t *Tensor) Dim(n int) int {
	return int(t.t.ne[n])
Michael Yang's avatar
Michael Yang committed
717
718
}

719
720
func (t *Tensor) Stride(n int) int {
	return int(t.t.nb[n])
Michael Yang's avatar
Michael Yang committed
721
722
}

723
724
func (t *Tensor) Shape() []int {
	shape := make([]int, C.ggml_n_dims(t.t))
Michael Yang's avatar
Michael Yang committed
725
726
727
728
729
730
731
	for i := range shape {
		shape[i] = t.Dim(i)
	}

	return shape
}

732
733
734
735
736
737
738
739
740
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
741
742
}

743
744
745
746
747
748
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
749
750
751
752
753
754
755
756
757
	}

	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
758
759
	case C.GGML_TYPE_F16:
		return ml.DTypeF16
760
761
762
763
	case C.GGML_TYPE_Q8_0:
		return ml.DTypeQ80
	case C.GGML_TYPE_Q4_0:
		return ml.DTypeQ40
Michael Yang's avatar
Michael Yang committed
764
765
766
767
768
769
770
	case C.GGML_TYPE_I32:
		return ml.DTypeI32
	default:
		return ml.DTypeOther
	}
}

771
772
773
774
775
776
777
func (t *Tensor) Neg(ctx ml.Context) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_neg(ctx.(*Context).ctx, t.t),
	}
}

Michael Yang's avatar
Michael Yang committed
778
779
func (t *Tensor) Add(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	return &Tensor{
780
		b: t.b,
Michael Yang's avatar
Michael Yang committed
781
782
783
784
		t: C.ggml_add(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
func (t *Tensor) Repeat(ctx ml.Context, dim, n int) ml.Tensor {
	if dim < 0 || dim >= C.GGML_MAX_DIMS {
		panic("invalid dimension")
	}

	shape := make([]C.int64_t, C.GGML_MAX_DIMS)
	for i := range C.GGML_MAX_DIMS {
		if i == dim {
			shape[i] = C.int64_t(t.Dim(i) * n)
		} else {
			shape[i] = C.int64_t(t.Dim(i))
		}
	}

	tmpl := C.ggml_new_tensor(ctx.(*Context).ctx, t.t._type, C.int(len(shape)), unsafe.SliceData(shape))
	return &Tensor{
		b: t.b,
		t: C.ggml_repeat(ctx.(*Context).ctx, t.t, tmpl),
	}
}

Michael Yang's avatar
Michael Yang committed
806
807
808
809
810
811
812
813
814
815
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{
816
		b: t.b,
Michael Yang's avatar
Michael Yang committed
817
818
819
820
821
822
		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{
823
		b: t.b,
Michael Yang's avatar
Michael Yang committed
824
825
826
827
828
829
		t: C.ggml_cont(ctx.(*Context).ctx, t.t),
	}
}

func (t *Tensor) Mul(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	return &Tensor{
830
		b: t.b,
Michael Yang's avatar
Michael Yang committed
831
832
833
834
835
836
		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{
837
		b: t.b,
Michael Yang's avatar
Michael Yang committed
838
839
840
841
		t: C.ggml_mul_mat(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

842
843
844
845
846
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{
847
		b: t.b,
848
849
850
851
		t: mul,
	}
}

Michael Yang's avatar
Michael Yang committed
852
func (t *Tensor) LayerNorm(ctx ml.Context, w, b ml.Tensor, eps float32) ml.Tensor {
853
	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
854
855
856
857
858
859
860
861
	if b != nil {
		tt = tt.Add(ctx, b)
	}

	return tt
}

func (t *Tensor) RMSNorm(ctx ml.Context, w ml.Tensor, eps float32) ml.Tensor {
862
	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
863
864
}

865
func (t *Tensor) Pad(ctx ml.Context, shape ...int) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
866
867
868
869
870
	if len(shape) != 4 {
		panic("expected 4 dimensions")
	}

	return &Tensor{
871
		b: t.b,
Michael Yang's avatar
Michael Yang committed
872
873
874
875
876
877
878
879
880
881
		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{
882
		b: t.b,
Michael Yang's avatar
Michael Yang committed
883
884
885
886
887
888
		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{
889
		b: t.b,
Michael Yang's avatar
Michael Yang committed
890
891
892
893
894
895
		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{
896
		b: t.b,
Michael Yang's avatar
Michael Yang committed
897
898
899
900
		t: C.ggml_cpy(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

901
func (t *Tensor) Reshape(ctx ml.Context, shape ...int) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
902
903
904
	switch len(shape) {
	case 1:
		return &Tensor{
905
			b: t.b,
Michael Yang's avatar
Michael Yang committed
906
907
908
909
			t: C.ggml_reshape_1d(ctx.(*Context).ctx, t.t, C.int64_t(shape[0])),
		}
	case 2:
		return &Tensor{
910
			b: t.b,
Michael Yang's avatar
Michael Yang committed
911
912
913
914
			t: C.ggml_reshape_2d(ctx.(*Context).ctx, t.t, C.int64_t(shape[0]), C.int64_t(shape[1])),
		}
	case 3:
		return &Tensor{
915
			b: t.b,
Michael Yang's avatar
Michael Yang committed
916
917
918
919
			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{
920
			b: t.b,
Michael Yang's avatar
Michael Yang committed
921
922
923
924
925
926
927
928
929
			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{
930
		b: t.b,
Michael Yang's avatar
Michael Yang committed
931
932
933
934
935
936
		t: C.ggml_scale(ctx.(*Context).ctx, t.t, (C.float)(s)),
	}
}

func (t *Tensor) Softmax(ctx ml.Context) ml.Tensor {
	return &Tensor{
937
		b: t.b,
Michael Yang's avatar
Michael Yang committed
938
939
940
941
		t: C.ggml_soft_max(ctx.(*Context).ctx, t.t),
	}
}

942
943
944
945
946
947
948
949
950
951
952
953
954
955
func (t *Tensor) Sin(ctx ml.Context) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_sin(ctx.(*Context).ctx, t.t),
	}
}

func (t *Tensor) Cos(ctx ml.Context) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_cos(ctx.(*Context).ctx, t.t),
	}
}

Michael Yang's avatar
Michael Yang committed
956
957
func (t *Tensor) Tanh(ctx ml.Context) ml.Tensor {
	return &Tensor{
958
		b: t.b,
Michael Yang's avatar
Michael Yang committed
959
960
961
962
		t: C.ggml_tanh_inplace(ctx.(*Context).ctx, t.t),
	}
}

963
func (t *Tensor) Unpad(ctx ml.Context, shape ...int) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
964
965
966
967
968
	if len(shape) != 4 {
		panic("expected 4 dimensions")
	}

	return &Tensor{
969
		b: t.b,
Michael Yang's avatar
Michael Yang committed
970
971
972
973
974
975
976
977
		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{
978
			b: t.b,
Michael Yang's avatar
Michael Yang committed
979
980
981
982
			t: C.ggml_view_1d(ctx.(*Context).ctx, t.t, C.int64_t(shape[0]), C.size_t(offset)),
		}
	case 3:
		return &Tensor{
983
			b: t.b,
Michael Yang's avatar
Michael Yang committed
984
985
986
987
988
989
990
			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{
991
			b: t.b,
Michael Yang's avatar
Michael Yang committed
992
993
994
995
996
997
998
			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{
999
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
			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 (
Patrick Devine's avatar
Patrick Devine committed
1011
1012
1013
1014
	ropeTypeNorm   C.int = 0
	ropeTypeNeox   C.int = 2
	ropeTypeMrope  C.int = 8
	ropeTypeVision C.int = 24
Michael Yang's avatar
Michael Yang committed
1015
1016
)

Patrick Devine's avatar
Patrick Devine committed
1017
func (t *Tensor) RoPE(ctx ml.Context, positionIDs, ropeFactors ml.Tensor, ropeDim, ropeType uint32, ropeBase, ropeScale float32) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
1018
	if ropeFactors == nil {
1019
		ropeFactors = &Tensor{b: t.b}
Michael Yang's avatar
Michael Yang committed
1020
1021
	}

Jesse Gross's avatar
Jesse Gross committed
1022
1023
1024
1025
1026
	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
1027
	return &Tensor{
1028
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1029
		t: C.ggml_rope_ext(
Jesse Gross's avatar
Jesse Gross committed
1030
			ctx.(*Context).ctx, dequant, positionIDs.(*Tensor).t, ropeFactors.(*Tensor).t,
Michael Yang's avatar
Michael Yang committed
1031
			C.int(ropeDim),
Patrick Devine's avatar
Patrick Devine committed
1032
1033
			C.int(ropeType),
			131072, // YaRN n_ctx_train
Michael Yang's avatar
Michael Yang committed
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
			C.float(ropeBase),
			C.float(ropeScale),
			0.,  // YaRN ext_factor
			1.,  // YaRN attn_factor
			32., // YaRN beta_fast
			1.,  // YaRN beta_slow
		),
	}
}

1044
1045
1046
1047
1048
1049
1050
func (t *Tensor) IM2Col(ctx ml.Context, t2 ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_im2col(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), true, C.GGML_TYPE_F32),
	}
}

Michael Yang's avatar
Michael Yang committed
1051
1052
func (t *Tensor) GELU(ctx ml.Context) ml.Tensor {
	return &Tensor{
1053
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1054
1055
1056
1057
1058
1059
		t: C.ggml_gelu_inplace(ctx.(*Context).ctx, t.t),
	}
}

func (t *Tensor) SILU(ctx ml.Context) ml.Tensor {
	return &Tensor{
1060
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1061
1062
1063
1064
1065
1066
		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{
1067
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1068
1069
1070
		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)),
	}
}
1071

Michael Yang's avatar
Michael Yang committed
1072
func (t *Tensor) AvgPool2D(ctx ml.Context, k, s int, p float32) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
1073
1074
	return &Tensor{
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1075
		t: C.ggml_pool_2d(ctx.(*Context).ctx, t.t, C.GGML_OP_POOL_AVG, C.int(k), C.int(k), C.int(s), C.int(s), C.float(p), C.float(p)),
Michael Yang's avatar
Michael Yang committed
1076
1077
1078
	}
}

Michael Yang's avatar
Michael Yang committed
1079
1080
1081
1082
func (t *Tensor) Set(ctx ml.Context, t2 ml.Tensor, offset int, strides ...int) ml.Tensor {
	var tt *C.struct_ggml_tensor
	switch len(strides) {
	case 0:
Michael Yang's avatar
Michael Yang committed
1083
		tt = C.ggml_set_1d(ctx.(*Context).ctx, t.t, t2.(*Tensor).t, C.size_t(offset))
Michael Yang's avatar
Michael Yang committed
1084
	case 1:
Michael Yang's avatar
Michael Yang committed
1085
		tt = C.ggml_set_2d(ctx.(*Context).ctx, t.t, t2.(*Tensor).t, C.size_t(offset), C.size_t(strides[0]))
Michael Yang's avatar
Michael Yang committed
1086
1087
1088
1089
1090
1091
1092
	default:
		panic("unsupported number of dimensions")
	}

	return &Tensor{b: t.b, t: tt}
}

1093
1094
1095
1096
1097
1098
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
	}

1099
1100
1101
	query := t.Permute(ctx, 0, 2, 1, 3)
	key = key.Permute(ctx, 0, 2, 1, 3)

1102
1103
	if t.b.flashAttention {
		value = value.Permute(ctx, 0, 2, 1, 3)
1104

1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
		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)
	}
1118
}
1119
1120
1121
1122
1123
1124
1125

func (t *Tensor) Duplicate(ctx ml.Context) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_dup(ctx.(*Context).ctx, t.t),
	}
}