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

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

Michael Yang's avatar
Michael Yang committed
36
37
38
39
40
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
41
	}
Michael Yang's avatar
Michael Yang committed
42
43

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

type Backend struct {
47
48
49
	// modelPath is the location of the model data
	modelPath string

50
51
	meta *fsggml.GGML

52
53
54
55
	// tensorLoadTargets maps from the name of the tensor in the file
	// to the name that is used by the model definition
	tensorLoadTargets map[string][]string

56
57
58
59
	sched         *C.struct_ggml_backend_sched
	schedBackends []*C.struct_ggml_backend
	schedBufts    []*C.struct_ggml_backend_buffer_type

60
	tensors map[string]*C.struct_ggml_tensor
Michael Yang's avatar
Michael Yang committed
61
62

	// input is the backend used for inputs
63
	input *C.struct_ggml_backend_buffer_type
Michael Yang's avatar
Michael Yang committed
64
65

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

68
69
70
71
72
73
	// requiredMemory is the cumulative memory allocations needed by the backend
	requiredMemory *ml.BackendMemory

	// btDeviceMemory maps from a buffer type to the memory allocations associated with that device
	btDeviceMemory map[*C.struct_ggml_backend_buffer_type]*ml.DeviceMemory

74
	flashAttention bool
Michael Yang's avatar
Michael Yang committed
75
76
77

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

80
81
82
83
84
85
86
87
func New(modelPath string, params ml.BackendParams) (ml.Backend, error) {
	r, err := os.Open(modelPath)
	if err != nil {
		return nil, err
	}
	defer r.Close()

	meta, err := fsggml.Decode(r, -1)
Michael Yang's avatar
Michael Yang committed
88
89
90
91
92
93
94
95
96
97
98
99
100
101
	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()),
	)

102
103
104
	var requiredMemory ml.BackendMemory
	btDeviceMemory := make(map[*C.struct_ggml_backend_buffer_type]*ml.DeviceMemory)

105
	type deviceBufferType struct {
106
107
108
109
110
		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
111
	for _, d := range devices() {
112
113
		switch C.ggml_backend_dev_type(d) {
		case C.GGML_BACKEND_DEVICE_TYPE_CPU:
114
115
116
117
			if len(cpus) == 0 {
				// only the first cpu device should be used
				cpus = append(cpus, d)
			}
118
119
		case C.GGML_BACKEND_DEVICE_TYPE_ACCEL:
			accels = append(accels, d)
Michael Yang's avatar
Michael Yang committed
120
		case C.GGML_BACKEND_DEVICE_TYPE_GPU:
121
			gpus = append(gpus, d)
Michael Yang's avatar
Michael Yang committed
122
123
124
		}
	}

125
126
	blocks := int(meta.KV().BlockCount())

Michael Yang's avatar
Michael Yang committed
127
	// create list of buffer types for the cpu
Michael Yang's avatar
Michael Yang committed
128
	cpuDeviceBufferType := deviceBufferType{d: C.ggml_backend_dev_by_type(C.GGML_BACKEND_DEVICE_TYPE_CPU)}
129
130
131
132
	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
133
			cpuDeviceBufferType.bts = append(cpuDeviceBufferType.bts, C.ggml_backend_dev_buffer_type(d))
134
			btDeviceMemory[C.ggml_backend_dev_buffer_type(d)] = &requiredMemory.CPU
Michael Yang's avatar
Michael Yang committed
135
		}
136
137
	}

138
	requiredMemory.CPU.Name = C.GoString(C.ggml_backend_dev_name(cpuDeviceBufferType.d))
139
140
	var props C.struct_ggml_backend_dev_props
	C.ggml_backend_dev_get_props(cpuDeviceBufferType.d, &props)
141
142
143
144

	// Bug #11211: Reporting of UUIDs is temporarily disabled due to causing segfaults
	// This only affects debug information until the new memory management code is in place
	// requiredMemory.CPU.UUID = C.GoString(props.uuid)
145
146
147
	requiredMemory.CPU.Weights = make([]ml.Memory, blocks+1)
	requiredMemory.CPU.Cache = make([]ml.Memory, blocks+1)

Michael Yang's avatar
Michael Yang committed
148
	// create list of buffer types for each gpu
149
	var gpuDeviceBufferTypes []deviceBufferType
150
151
	requiredMemory.GPUs = make([]ml.DeviceMemory, len(gpus))
	for i, d := range gpus {
152
		bt := C.ggml_backend_dev_buffer_type(d)
153
		gpuDeviceBufferTypes = append(gpuDeviceBufferTypes, deviceBufferType{
154
			d:   d,
Michael Yang's avatar
Michael Yang committed
155
			bts: append([]*C.struct_ggml_backend_buffer_type{bt}, cpuDeviceBufferType.bts...),
156
		})
157
158
		btDeviceMemory[bt] = &requiredMemory.GPUs[i]
		requiredMemory.GPUs[i].Name = C.GoString(C.ggml_backend_dev_name(d))
159
160
		var props C.struct_ggml_backend_dev_props
		C.ggml_backend_dev_get_props(d, &props)
161
		// requiredMemory.GPUs[i].UUID = C.GoString(props.uuid)
162
163
		requiredMemory.GPUs[i].Weights = make([]ml.Memory, blocks+1)
		requiredMemory.GPUs[i].Cache = make([]ml.Memory, blocks+1)
Michael Yang's avatar
Michael Yang committed
164
165
	}

Michael Yang's avatar
Michael Yang committed
166
167
168
169
170
	useDefaultSplit := true
	for _, s := range params.TensorSplit {
		if s != 0 {
			useDefaultSplit = false
			break
171
		}
Michael Yang's avatar
Michael Yang committed
172
	}
173

Michael Yang's avatar
Michael Yang committed
174
175
176
177
	// calculate splits
	splits := make([]float32, len(gpus))
	if useDefaultSplit {
		// default: split on free memory
178
179
180
181
182
		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
183
184
	} else {
		splits = params.TensorSplit
185
186
187
	}

	var sum float32
Michael Yang's avatar
Michael Yang committed
188
	// cumulative sum of all splits
189
190
191
192
193
	for i := range splits {
		sum += splits[i]
		splits[i] = sum
	}

Michael Yang's avatar
Michael Yang committed
194
	// normalize splits
195
	for i := range splits {
196
		splits[i] /= sum
197
198
	}

Michael Yang's avatar
Michael Yang committed
199
	// inputs always use cpu
Michael Yang's avatar
Michael Yang committed
200
	input := cpuDeviceBufferType
201

Michael Yang's avatar
Michael Yang committed
202
203
204
	// 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
205
	assignLayer := func(i int) deviceBufferType {
Michael Yang's avatar
Michael Yang committed
206
		if i < gpuRangeStart || i >= gpuRangeStop {
Michael Yang's avatar
Michael Yang committed
207
			return cpuDeviceBufferType
208
		}
209

Michael Yang's avatar
Michael Yang committed
210
		index := slices.IndexFunc(splits, func(f float32) bool { return float32(i-gpuRangeStart)/float32(gpuRangeStop-gpuRangeStart) < f })
211
		if index < 0 || index >= len(gpuDeviceBufferTypes) {
Michael Yang's avatar
Michael Yang committed
212
			return cpuDeviceBufferType
213
214
215
		}

		return gpuDeviceBufferTypes[index]
216
217
	}

Michael Yang's avatar
Michael Yang committed
218
	// repeating layers are assigned based on their index in reverse order, e.g. i / (block_count + 1)
219
	layers := make([]deviceBufferType, blocks)
220
	for i := range layers {
221
		layers[i] = assignLayer(i)
222
223
	}

Michael Yang's avatar
Michael Yang committed
224
	// outputs are assigned iff allowed by splits and configured number of gpu layers
225
	output := assignLayer(blocks)
226
227
228

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

232
	type tensor struct {
233
		source *fsggml.Tensor
234
235
236
		target string
	}

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

Michael Yang's avatar
Michael Yang committed
240
	// contexts are shared by tensors of the same buffer type
241
	ctxs := make(map[*C.struct_ggml_backend_buffer_type]*C.struct_ggml_context)
242
	createTensor := func(t tensor, bts []*C.struct_ggml_backend_buffer_type, layer int) *C.struct_ggml_tensor {
243
244
245
246
247
248
249
		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
250

251
252
253
254
255
256
257
258
			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
259
			defer C.free(unsafe.Pointer(cname))
260
261
262
263
			if tt := C.ggml_get_tensor(ctxs[bt], cname); tt != nil {
				return tt
			}

264
			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
265
266
			C.ggml_set_name(tt, cname)

267
			slog.Log(context.TODO(), logutil.LevelTrace, "created tensor", "name", name, "shape", t.source.Shape, "dtype", t.source.Kind, "buffer_type", C.GoString(C.ggml_backend_buft_name(bt)))
268
269
270
271
272
273
274
275
276
277

			size := pad(C.ggml_backend_buft_get_alloc_size(bt, tt), C.ggml_backend_buft_get_alignment(bt))
			if layer == -1 {
				// Assume that InputWeights can be allocated - they're always in system memory and can't be moved in any case
				requiredMemory.InputWeights.Status = ml.Allocated
				requiredMemory.InputWeights.Size += uint64(size)
			} else {
				btDeviceMemory[bt].Weights[layer].Size += uint64(size)
			}

278
279
280
281
282
			//nolint:staticcheck // TODO: check if buffer type supports this tensor
			return tt
		}

		return nil
Michael Yang's avatar
Michael Yang committed
283
284
	}

285
	contains := func(s string, parts ...string) bool {
286
287
288
289
290
291
292
293
		split := strings.Split(s, ".")
		for _, part := range parts {
			if slices.Contains(split, part) {
				return true
			}
		}

		return false
Michael Yang's avatar
Michael Yang committed
294
295
	}

296
297
	for _, t := range meta.Tensors().Items() {
		switch {
298
		case contains(t.Name, "position_embd", "token_embd", "token_norm_embd", "token_types"):
299
			createTensor(tensor{source: t}, input.bts, -1)
Michael Yang's avatar
Michael Yang committed
300
			if _, ok := meta.Tensors().GroupLayers()["output"]; !ok && t.Name == "token_embd.weight" {
301
				createTensor(tensor{source: t, target: "output.weight"}, output.bts, blocks)
Michael Yang's avatar
Michael Yang committed
302
			}
Michael Yang's avatar
Michael Yang committed
303
304
305
		case contains(t.Name, "cls", "output", "output_norm",
			"altup_proj", "altup_unembd_proj",
			"per_layer_token_embd", "per_layer_model_proj", "per_layer_proj_norm"):
306
			createTensor(tensor{source: t}, output.bts, blocks)
307
		case strings.HasPrefix(t.Name, "v.") || strings.HasPrefix(t.Name, "mm."):
Michael Yang's avatar
Michael Yang committed
308
			// TODO: assign vision tensors to the gpu if possible
309
			createTensor(tensor{source: t}, output.bts, blocks)
Michael Yang's avatar
Michael Yang committed
310
311
312
313
314
315
		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,
316
				}, layer.bts, i)
Michael Yang's avatar
Michael Yang committed
317
			}
318
		default:
Michael Yang's avatar
Michael Yang committed
319
320
321
322
			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
323
				}
Michael Yang's avatar
Michael Yang committed
324
			}
325

Michael Yang's avatar
Michael Yang committed
326
			if layerIndex >= 0 {
327
				createTensor(tensor{source: t}, layers[layerIndex].bts, layerIndex)
328
			} else {
Michael Yang's avatar
Michael Yang committed
329
				// load all other tensors on the cpu
330
				createTensor(tensor{source: t}, input.bts, -1)
331
332
333
			}
		}
	}
Michael Yang's avatar
Michael Yang committed
334

Michael Yang's avatar
Michael Yang committed
335
336
	// allocate buffers for each context
	bbs := make(map[*C.struct_ggml_context]*C.struct_ggml_backend_buffer, len(ctxs))
337
338
339
340
341
342
	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)
343
344
345
346
347
348
349
350
351
352
		for i := range btDeviceMemory[bt].Weights {
			if btDeviceMemory[bt].Weights[i].Size != 0 {
				if b != nil {
					btDeviceMemory[bt].Weights[i].Status = ml.Allocated
				} else {
					btDeviceMemory[bt].Weights[i].Status = ml.Failed
				}
			}
		}

353
		if b == nil {
354
			panic(ml.ErrNoMem{BackendMemory: requiredMemory})
355
356
		}

357
		C.ggml_backend_buffer_set_usage(b, C.GGML_BACKEND_BUFFER_USAGE_WEIGHTS)
Michael Yang's avatar
Michael Yang committed
358
		bbs[c] = b
359
360
361
	}

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

Michael Yang's avatar
Michael Yang committed
365
	// map tensor names to tensors for easy lookup later
366
367
368
369
370
371
372
	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
		}
	}

373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
	// 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)

	// 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
	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)

		deviceBufferTypes[d] = bt

		schedBackends = append(schedBackends, b)
		schedBufts = append(schedBufts, bt)

		if C.ggml_backend_is_cpu(b) {
			// set number of threads for cpu backend
			C.ggml_backend_cpu_set_n_threads(b, C.int(Threads(params.NumThreads)))
		}
	}

	maxGraphNodes := max(8192, len(meta.Tensors().Items())*5)
	return &Backend{
		modelPath:         modelPath,
		flashAttention:    params.FlashAttention,
		meta:              meta,
		tensorLoadTargets: targets,
		tensors:           tensors,
		sched: C.ggml_backend_sched_new(
			(*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),
			C._Bool(len(gpus) > 1 && slices.Contains(gpus, output.d)),
			C._Bool(false),
		),
		schedBackends: schedBackends,
		schedBufts:    schedBufts,
		input:         deviceBufferTypes[input.d],
		layers: func() map[int]*C.struct_ggml_backend_buffer_type {
			m := make(map[int]*C.struct_ggml_backend_buffer_type)
			for i, layer := range layers {
				m[i] = deviceBufferTypes[layer.d]
			}
			return m
		}(),
419
420
421
		requiredMemory: &requiredMemory,
		btDeviceMemory: btDeviceMemory,
		maxGraphNodes:  maxGraphNodes,
422
423
424
425
426
427
428
429
	}, nil
}

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

func (b *Backend) Load(ctx context.Context, progress func(float32)) error {
430
	var doneBytes atomic.Uint64
431
	totalBytes := uint64(b.meta.Length) - b.meta.Tensors().Offset
432
433
434

	g, ctx := errgroup.WithContext(ctx)
	g.SetLimit(runtime.GOMAXPROCS(0))
435
	for _, t := range b.meta.Tensors().Items() {
436
		t := t
437
		g.Go(func() error {
438
			tts := make([]*C.struct_ggml_tensor, max(1, len(b.tensorLoadTargets[t.Name])))
439
			for i := range tts {
440
				target := b.tensorLoadTargets[t.Name][i]
441
442
443
				if target == "" {
					target = t.Name
				}
444

445
				tt, ok := b.tensors[target]
446
447
448
				if !ok {
					return fmt.Errorf("unassigned tensor: %s", t.Name)
				}
Michael Yang's avatar
Michael Yang committed
449

450
451
452
				tts[i] = tt
			}

453
454
			// 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.
455
			file, err := os.Open(b.modelPath)
456
			if err != nil {
457
				slog.Warn("file open error", "file", b.modelPath, "error", err)
458
459
460
				return err
			}
			defer file.Close()
461
			sr := io.NewSectionReader(file, int64(b.meta.Tensors().Offset+t.Offset), int64(t.Size()))
462
463
464
465
			bts := make([]byte, 128*format.KibiByte)

			var s uint64
			for s < t.Size() {
466
467
468
469
470
				// Stop if either the parent context has been canceled or if any of the other tensors returned an error
				if err := ctx.Err(); err != nil {
					return err
				}

471
472
				n, err := io.ReadFull(sr, bts[:min(len(bts), int(t.Size()-s))])
				if err != nil {
473
					slog.Warn("file read error", "file", b.modelPath, "error", err)
474
					return err
475
				}
Michael Yang's avatar
Michael Yang committed
476

477
478
				for _, tt := range tts {
					C.ggml_backend_tensor_set(tt, unsafe.Pointer(&bts[0]), C.size_t(s), C.size_t(n))
479
				}
Michael Yang's avatar
Michael Yang committed
480

481
482
				s += uint64(n)

483
				if progress != nil {
484
					done := doneBytes.Add(uint64(n))
485
					progress(float32(done) / float32(totalBytes))
486
487
488
489
490
				}
			}

			return nil
		})
Michael Yang's avatar
Michael Yang committed
491
492
	}

493
	if err := g.Wait(); err != nil {
494
		return err
495
496
	}

497
	return nil
Michael Yang's avatar
Michael Yang committed
498
499
}

500
501
502
503
func (b *Backend) BackendMemory() ml.BackendMemory {
	return *b.requiredMemory
}

504
func (b *Backend) Config() fs.Config {
Michael Yang's avatar
Michael Yang committed
505
506
507
508
	return b.meta.KV()
}

func (b *Backend) Get(name string) ml.Tensor {
509
510
	if t, ok := b.tensors[name]; ok {
		return &Tensor{b: b, t: t}
Michael Yang's avatar
Michael Yang committed
511
512
513
514
515
516
	}

	return nil
}

func (b *Backend) NewContext() ml.Context {
Michael Yang's avatar
Michael Yang committed
517
	return b.NewContextSize(b.maxGraphNodes)
518
519
520
}

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

525
526
	var allocatedBuffers []*C.struct_ggml_backend_buffer

Michael Yang's avatar
Michael Yang committed
527
	return &Context{
528
529
		b:             b,
		maxGraphNodes: n,
530
		ctx: C.ggml_init(C.struct_ggml_init_params{
531
			mem_size: C.size_t(n)*C.ggml_tensor_overhead() + C.ggml_graph_overhead_custom(C.size_t(n), false),
532
533
			no_alloc: true,
		}),
534
		allocatedBuffers: &allocatedBuffers,
535
		layer:            -1,
Michael Yang's avatar
Michael Yang committed
536
537
538
	}
}

539
func (b *Backend) CacheConfig() ml.CacheConfig {
540
541
542
543
544
	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}
	}
545
546
}

Michael Yang's avatar
Michael Yang committed
547
type Context struct {
548
	b *Backend
Michael Yang's avatar
Michael Yang committed
549

550
	ctx   *C.struct_ggml_context
Michael Yang's avatar
Michael Yang committed
551
	graph *C.struct_ggml_cgraph
552

553
554
	// buft is the buffer type used for new tensors
	buft *C.struct_ggml_backend_buffer_type
555

556
557
558
559
	// allocatedBuffers are buffers for tensors that we have allocated in this context
	// so that we can free them when we close the context
	allocatedBuffers *[]*C.struct_ggml_backend_buffer

Michael Yang's avatar
Michael Yang committed
560
	// maxGraphNodes is the maximum allowed number of graph nodes in this context
561
	maxGraphNodes int
562
563
564

	// layer is the graph layer that this context is allocating for - assumed to be cache
	layer int
Michael Yang's avatar
Michael Yang committed
565
566
}

567
func (c *Context) Input() ml.Context {
Michael Yang's avatar
Michael Yang committed
568
	if c.b.input != nil {
569
		return &Context{
570
571
572
573
574
			b:                c.b,
			ctx:              c.ctx,
			buft:             c.b.input,
			allocatedBuffers: c.allocatedBuffers,
			maxGraphNodes:    c.maxGraphNodes,
575
			layer:            -1,
576
577
578
		}
	}

579
	return c
580
581
}

582
func (c *Context) Layer(i int) ml.Context {
583
	if buft, ok := c.b.layers[i]; ok {
584
		return &Context{
585
586
587
588
589
			b:                c.b,
			ctx:              c.ctx,
			buft:             buft,
			allocatedBuffers: c.allocatedBuffers,
			maxGraphNodes:    c.maxGraphNodes,
590
			layer:            i,
591
592
593
		}
	}

594
	return c
595
596
}

597
func (c *Context) Forward(tensors ...ml.Tensor) ml.Context {
Michael Yang's avatar
Michael Yang committed
598
	if c.graph == nil {
599
		c.graph = C.ggml_new_graph_custom(c.ctx, C.size_t(c.maxGraphNodes), false)
Michael Yang's avatar
Michael Yang committed
600
601
	}

602
603
604
605
606
	for _, tensor := range tensors {
		C.ggml_build_forward_expand(c.graph, tensor.(*Tensor).t)
	}

	return c
Michael Yang's avatar
Michael Yang committed
607
608
}

609
func (c *Context) Compute(tensors ...ml.Tensor) {
610
611
612
	if status := C.ggml_backend_sched_graph_compute_async(c.b.sched, c.graph); status != C.GGML_STATUS_SUCCESS {
		panic(fmt.Errorf("error computing ggml graph: %v", status))
	}
Michael Yang's avatar
Michael Yang committed
613
	C.ggml_backend_sched_reset(c.b.sched)
Michael Yang's avatar
Michael Yang committed
614

615
616
617
	needSync := true
	sync := func() {
		if needSync {
618
			C.ggml_backend_sched_synchronize(c.b.sched)
619
620
621
			needSync = false
		}
	}
Michael Yang's avatar
Michael Yang committed
622

623
624
625
	for _, t := range tensors {
		if C.ggml_nbytes(t.(*Tensor).t) > 0 {
			t.(*Tensor).sync = sync
626
627
		}
	}
Michael Yang's avatar
Michael Yang committed
628
629
}

630
631
func (c *Context) Reserve() {
	reserved := C.ggml_backend_sched_reserve(c.b.sched, c.graph)
632
633

	slog.Debug("compute graph", "nodes", C.ggml_graph_n_nodes(c.graph), "splits", C.ggml_backend_sched_get_n_splits(c.b.sched))
634
635
636
637
638
639

	// Reserve may get called multiple times for different graphs - we just want the last run, which will contain the max allocations
	for _, bt := range c.b.schedBufts {
		c.b.btDeviceMemory[bt].Graph = ml.Memory{}
	}

640
	for i := range c.b.schedBackends {
641
642
643
644
645
646
647
648
649
650
		bufferStatus := C.ggml_backend_sched_get_attempted_buffer_size(c.b.sched, c.b.schedBackends[i])

		graph := &c.b.btDeviceMemory[c.b.schedBufts[i]].Graph
		graph.Size += uint64(bufferStatus.size)
		if bufferStatus.allocated && graph.Status != ml.Failed {
			graph.Status = ml.Allocated
		} else {
			graph.Status = ml.Failed
		}

651
		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])),
652
			"size", format.HumanBytes2(uint64(bufferStatus.size)))
653
654
	}

655
656
657
	if !reserved {
		panic(ml.ErrNoMem{BackendMemory: *c.b.requiredMemory})
	}
658
659
}

660
func (c *Context) MaxGraphNodes() int {
661
	return c.maxGraphNodes
Jesse Gross's avatar
Jesse Gross committed
662
663
}

664
665
666
func shapeToGGML(shape []int) *C.int64_t {
	sh := make([]C.int64_t, len(shape))
	for i, s := range shape {
667
		sh[i] = C.int64_t(s)
668
669
670
671
672
	}

	return &sh[0]
}

673
674
675
676
func pad(length, pad C.size_t) C.size_t {
	return ((length + pad - 1) / pad) * pad
}

677
func (c *Context) newTensor(dtype ml.DType, shape []int) ml.Tensor {
678
	if c.buft == nil {
679
		panic("set Input or Layer before creating tensors")
680
681
	}

Michael Yang's avatar
Michael Yang committed
682
683
684
685
686
687
	var cdtype uint32
	switch dtype {
	case ml.DTypeF32:
		cdtype = C.GGML_TYPE_F32
	case ml.DTypeF16:
		cdtype = C.GGML_TYPE_F16
688
689
690
691
	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
692
693
694
695
696
697
	case ml.DTypeI32:
		cdtype = C.GGML_TYPE_I32
	default:
		panic("unsupported dtype")
	}

Jesse Gross's avatar
Jesse Gross committed
698
	if len(shape) < 1 || shape[0] == 0 {
Michael Yang's avatar
Michael Yang committed
699
		var shape C.int64_t = 0
700
		return &Tensor{b: c.b, t: C.ggml_new_tensor(c.ctx, cdtype, 1, &shape)}
Michael Yang's avatar
Michael Yang committed
701
	} else if len(shape) > 4 {
Michael Yang's avatar
Michael Yang committed
702
703
704
705
706
707
708
709
710
		panic("unsupported number of dimensions")
	}

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

Michael Yang's avatar
Michael Yang committed
711
	t := C.ggml_new_tensor(c.ctx, cdtype, C.int(len(shape)), shapeToGGML(shape))
712
	size := pad(C.ggml_backend_buft_get_alloc_size(c.buft, t), C.ggml_backend_buft_get_alignment(c.buft))
713

714
	b := C.ggml_backend_buft_alloc_buffer(c.buft, size)
715
716
717
718
719
720
721
722
723
724
725
	if c.layer >= 0 {
		cache := &c.b.btDeviceMemory[c.buft].Cache[c.layer]

		cache.Size += uint64(size)
		if b != nil {
			cache.Status = ml.Allocated
		} else {
			cache.Status = ml.Failed
		}
	}

726
	if b == nil {
727
		panic(ml.ErrNoMem{BackendMemory: *c.b.requiredMemory})
728
729
	}

730
	*c.allocatedBuffers = append(*c.allocatedBuffers, b)
Michael Yang's avatar
Michael Yang committed
731
	C.ggml_backend_tensor_alloc(b, t, C.ggml_backend_buffer_get_base(b))
732
	return &Tensor{b: c.b, t: t}
733
734
}

735
func (c *Context) Empty(dtype ml.DType, shape ...int) ml.Tensor {
736
	return c.newTensor(dtype, shape)
737
738
}

739
func (c *Context) Zeros(dtype ml.DType, shape ...int) ml.Tensor {
740
	t := c.newTensor(dtype, shape)
741
742
	C.ggml_set_zero(t.(*Tensor).t)
	return t
Michael Yang's avatar
Michael Yang committed
743
744
}

745
func checkShape[S ~[]E, E any](s S, shape ...int) {
Michael Yang's avatar
Michael Yang committed
746
	n := len(s)
Jesse Gross's avatar
Jesse Gross committed
747
748

	if n == 0 {
749
		return
Jesse Gross's avatar
Jesse Gross committed
750
751
	}

Michael Yang's avatar
Michael Yang committed
752
753
754
755
756
	for _, v := range shape {
		n /= v
	}

	if n != 1 {
757
		panic(fmt.Errorf("invalid shape: %v", shape))
Michael Yang's avatar
Michael Yang committed
758
759
760
	}
}

761
762
func (c *Context) FromFloatSlice(s []float32, shape ...int) ml.Tensor {
	checkShape(s, shape...)
763

764
	t := c.newTensor(ml.DTypeF32, shape)
765

Jesse Gross's avatar
Jesse Gross committed
766
767
768
769
	if len(s) > 0 {
		C.ggml_backend_tensor_set(t.(*Tensor).t, unsafe.Pointer(&s[0]), 0, C.ggml_nbytes(t.(*Tensor).t))
	}

770
	return t
Michael Yang's avatar
Michael Yang committed
771
772
}

773
774
func (c *Context) FromIntSlice(s []int32, shape ...int) ml.Tensor {
	checkShape(s, shape...)
775

776
	t := c.newTensor(ml.DTypeI32, shape)
777

Jesse Gross's avatar
Jesse Gross committed
778
779
780
781
	if len(s) > 0 {
		C.ggml_backend_tensor_set(t.(*Tensor).t, unsafe.Pointer(&s[0]), 0, C.ggml_nbytes(t.(*Tensor).t))
	}

782
	return t
Michael Yang's avatar
Michael Yang committed
783
784
}

Michael Yang's avatar
arange  
Michael Yang committed
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
func (c Context) Arange(start, stop, step float32, dtype ml.DType) ml.Tensor {
	switch dtype {
	case ml.DTypeF32:
		// ggml_arange creates a float32 tensor
		return &Tensor{
			b: c.b,
			t: C.ggml_arange(c.ctx, C.float(start), C.float(stop), C.float(step)),
		}
	case ml.DTypeI32:
		// ggml_cast does not support float32 to int32 conversion
		arange := make([]int32, 0, int((stop-start)/step))
		for i := start; i < stop; i += step {
			arange = append(arange, int32(i))
		}

800
		return c.Input().FromIntSlice(arange, len(arange))
Michael Yang's avatar
arange  
Michael Yang committed
801
802
803
804
805
	default:
		panic("unsupported dtype for arange")
	}
}

Michael Yang's avatar
Michael Yang committed
806
807
func (c *Context) Close() {
	if c != nil {
808
809
810
811
812
		for _, b := range *c.allocatedBuffers {
			C.ggml_backend_buffer_free(b)
		}
		*c.allocatedBuffers = nil

813
814
		C.ggml_free(c.ctx)
	}
Michael Yang's avatar
Michael Yang committed
815
816
817
}

type Tensor struct {
818
	b    *Backend
Michael Yang's avatar
Michael Yang committed
819
	t    *C.struct_ggml_tensor
820
	sync func()
Michael Yang's avatar
Michael Yang committed
821
822
823
824
825
826
827
828
829
830
}

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

831
832
func (t *Tensor) Dim(n int) int {
	return int(t.t.ne[n])
Michael Yang's avatar
Michael Yang committed
833
834
}

835
836
func (t *Tensor) Stride(n int) int {
	return int(t.t.nb[n])
Michael Yang's avatar
Michael Yang committed
837
838
}

839
840
func (t *Tensor) Shape() []int {
	shape := make([]int, C.ggml_n_dims(t.t))
Michael Yang's avatar
Michael Yang committed
841
842
843
844
845
846
847
	for i := range shape {
		shape[i] = t.Dim(i)
	}

	return shape
}

848
849
850
851
852
853
854
855
856
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
857
858
}

859
860
861
862
863
864
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
865
866
867
868
869
870
871
872
873
	}

	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
874
875
	case C.GGML_TYPE_F16:
		return ml.DTypeF16
876
877
878
879
	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
880
881
882
883
884
885
886
	case C.GGML_TYPE_I32:
		return ml.DTypeI32
	default:
		return ml.DTypeOther
	}
}

887
888
889
890
891
892
893
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
894
895
func (t *Tensor) Add(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_add(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

Michael Yang's avatar
Michael Yang committed
901
902
903
904
905
906
907
func (t *Tensor) Sub(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_sub(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
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
929
930
931
932
933
934
935
936
937
938
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{
939
		b: t.b,
Michael Yang's avatar
Michael Yang committed
940
941
942
943
944
945
		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{
946
		b: t.b,
Michael Yang's avatar
Michael Yang committed
947
948
949
950
951
952
		t: C.ggml_cont(ctx.(*Context).ctx, t.t),
	}
}

func (t *Tensor) Mul(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	return &Tensor{
953
		b: t.b,
Michael Yang's avatar
Michael Yang committed
954
955
956
957
		t: C.ggml_mul(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

958
959
960
961
962
963
964
func (t *Tensor) Div(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_div(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

Michael Yang's avatar
Michael Yang committed
965
966
func (t *Tensor) Mulmat(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	return &Tensor{
967
		b: t.b,
Michael Yang's avatar
Michael Yang committed
968
969
970
971
		t: C.ggml_mul_mat(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

972
973
974
975
976
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{
977
		b: t.b,
978
979
980
981
		t: mul,
	}
}

Michael Yang's avatar
llama4  
Michael Yang committed
982
983
984
985
986
987
988
func (t *Tensor) MulmatID(ctx ml.Context, t2, ids ml.Tensor) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_mul_mat_id(ctx.(*Context).ctx, t.t, t2.(*Tensor).t, ids.(*Tensor).t),
	}
}

Michael Yang's avatar
Michael Yang committed
989
func (t *Tensor) LayerNorm(ctx ml.Context, w, b ml.Tensor, eps float32) ml.Tensor {
Michael Yang's avatar
llama4  
Michael Yang committed
990
991
992
993
994
995
	tt := C.ggml_norm(ctx.(*Context).ctx, t.t, C.float(eps))
	if w != nil {
		tt = C.ggml_mul(ctx.(*Context).ctx, tt, w.(*Tensor).t)
		if b != nil {
			tt = C.ggml_add(ctx.(*Context).ctx, tt, b.(*Tensor).t)
		}
Michael Yang's avatar
Michael Yang committed
996
997
	}

Michael Yang's avatar
llama4  
Michael Yang committed
998
	return &Tensor{b: t.b, t: tt}
Michael Yang's avatar
Michael Yang committed
999
1000
1001
}

func (t *Tensor) RMSNorm(ctx ml.Context, w ml.Tensor, eps float32) ml.Tensor {
Michael Yang's avatar
llama4  
Michael Yang committed
1002
1003
1004
1005
1006
1007
	tt := C.ggml_rms_norm(ctx.(*Context).ctx, t.t, C.float(eps))
	if w != nil {
		tt = C.ggml_mul(ctx.(*Context).ctx, tt, w.(*Tensor).t)
	}

	return &Tensor{b: t.b, t: tt}
Michael Yang's avatar
Michael Yang committed
1008
1009
}

1010
func (t *Tensor) Pad(ctx ml.Context, shape ...int) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
1011
1012
	if len(shape) != 4 {
		panic("expected 4 dimensions")
1013
1014
	} else if shape[3] != 0 {
		panic("cuda does not support 4d tensors")
Michael Yang's avatar
Michael Yang committed
1015
1016
1017
	}

	return &Tensor{
1018
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
		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{
1029
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1030
1031
1032
1033
1034
1035
		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{
1036
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1037
1038
1039
1040
1041
1042
		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{
1043
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1044
1045
1046
1047
		t: C.ggml_cpy(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

1048
func (t *Tensor) Reshape(ctx ml.Context, shape ...int) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
1049
1050
1051
	switch len(shape) {
	case 1:
		return &Tensor{
1052
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1053
1054
1055
1056
			t: C.ggml_reshape_1d(ctx.(*Context).ctx, t.t, C.int64_t(shape[0])),
		}
	case 2:
		return &Tensor{
1057
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1058
1059
1060
1061
			t: C.ggml_reshape_2d(ctx.(*Context).ctx, t.t, C.int64_t(shape[0]), C.int64_t(shape[1])),
		}
	case 3:
		return &Tensor{
1062
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1063
1064
1065
1066
			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{
1067
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1068
1069
1070
1071
1072
1073
1074
1075
1076
			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{
1077
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1078
1079
1080
1081
		t: C.ggml_scale(ctx.(*Context).ctx, t.t, (C.float)(s)),
	}
}

1082
1083
1084
1085
1086
1087
1088
func (t *Tensor) SumRows(ctx ml.Context) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_sum_rows(ctx.(*Context).ctx, t.t),
	}
}

Michael Yang's avatar
Michael Yang committed
1089
1090
func (t *Tensor) Softmax(ctx ml.Context) ml.Tensor {
	return &Tensor{
1091
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1092
1093
1094
1095
		t: C.ggml_soft_max(ctx.(*Context).ctx, t.t),
	}
}

1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
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
1110
1111
func (t *Tensor) Tanh(ctx ml.Context) ml.Tensor {
	return &Tensor{
1112
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1113
1114
1115
1116
		t: C.ggml_tanh_inplace(ctx.(*Context).ctx, t.t),
	}
}

Michael Yang's avatar
llama4  
Michael Yang committed
1117
1118
1119
1120
1121
1122
1123
func (t *Tensor) Sigmoid(ctx ml.Context) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_sigmoid_inplace(ctx.(*Context).ctx, t.t),
	}
}

Michael Yang's avatar
Michael Yang committed
1124
1125
1126
1127
func (t *Tensor) View(ctx ml.Context, offset int, shape ...int) ml.Tensor {
	switch len(shape) {
	case 1:
		return &Tensor{
1128
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1129
1130
1131
1132
			t: C.ggml_view_1d(ctx.(*Context).ctx, t.t, C.int64_t(shape[0]), C.size_t(offset)),
		}
	case 3:
		return &Tensor{
1133
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1134
1135
1136
1137
1138
1139
1140
			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{
1141
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1142
1143
1144
1145
1146
1147
1148
			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{
1149
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
			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")
	}
}

1160
func (t *Tensor) RoPE(ctx ml.Context, positions ml.Tensor, ropeDim int, ropeBase, ropeScale float32, options ...func(*rope.Options)) ml.Tensor {
1161
	// Default options
1162
	opts := &rope.Options{OriginalContextLength: 131072, Factors: &Tensor{}}
1163
1164
1165
1166
1167
1168

	// Apply any provided options
	for _, option := range options {
		option(opts)
	}

Jesse Gross's avatar
Jesse Gross committed
1169
1170
1171
1172
1173
	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
1174
	return &Tensor{
1175
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1176
		t: C.ggml_rope_ext(
1177
1178
			ctx.(*Context).ctx,
			dequant,
1179
1180
			positions.(*Tensor).t,
			opts.Factors.(*Tensor).t,
Michael Yang's avatar
Michael Yang committed
1181
			C.int(ropeDim),
1182
1183
			C.int(opts.Type),
			C.int(opts.OriginalContextLength),
Michael Yang's avatar
Michael Yang committed
1184
1185
			C.float(ropeBase),
			C.float(ropeScale),
1186
1187
1188
1189
			C.float(0.0),
			C.float(1.0),
			C.float(32.0),
			C.float(1.0),
Michael Yang's avatar
Michael Yang committed
1190
1191
1192
1193
		),
	}
}

1194
1195
1196
1197
1198
1199
1200
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
1201
1202
func (t *Tensor) GELU(ctx ml.Context) ml.Tensor {
	return &Tensor{
1203
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1204
1205
1206
1207
1208
1209
		t: C.ggml_gelu_inplace(ctx.(*Context).ctx, t.t),
	}
}

func (t *Tensor) SILU(ctx ml.Context) ml.Tensor {
	return &Tensor{
1210
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1211
1212
1213
1214
		t: C.ggml_silu_inplace(ctx.(*Context).ctx, t.t),
	}
}

Michael Yang's avatar
Michael Yang committed
1215
1216
1217
1218
1219
1220
1221
func (t *Tensor) RELU(ctx ml.Context) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_relu_inplace(ctx.(*Context).ctx, t.t),
	}
}

Michael Yang's avatar
Michael Yang committed
1222
1223
func (t *Tensor) Conv2D(ctx ml.Context, t2 ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor {
	return &Tensor{
1224
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1225
1226
1227
		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)),
	}
}
1228

Michael Yang's avatar
Michael Yang committed
1229
func (t *Tensor) AvgPool2D(ctx ml.Context, k, s int, p float32) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
1230
1231
	return &Tensor{
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1232
		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
1233
1234
1235
	}
}

Michael Yang's avatar
Michael Yang committed
1236
1237
1238
1239
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
1240
		tt = C.ggml_set_1d(ctx.(*Context).ctx, t.t, t2.(*Tensor).t, C.size_t(offset))
Michael Yang's avatar
Michael Yang committed
1241
	case 1:
Michael Yang's avatar
Michael Yang committed
1242
		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
1243
1244
1245
1246
1247
1248
1249
	default:
		panic("unsupported number of dimensions")
	}

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

1250
1251
1252
1253
1254
1255
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
	}

1256
1257
1258
	query := t.Permute(ctx, 0, 2, 1, 3)
	key = key.Permute(ctx, 0, 2, 1, 3)

1259
1260
	if t.b.flashAttention {
		value = value.Permute(ctx, 0, 2, 1, 3)
1261

1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
		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)
	}
1275
}
1276
1277
1278
1279
1280
1281
1282

func (t *Tensor) Duplicate(ctx ml.Context) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_dup(ctx.(*Context).ctx, t.t),
	}
}
Michael Yang's avatar
llama4  
Michael Yang committed
1283
1284
1285
1286
1287
1288
1289

func (t *Tensor) TopK(ctx ml.Context, k int) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_top_k(ctx.(*Context).ctx, t.t, C.int(k)),
	}
}
1290
1291
1292
1293
1294
1295
1296

func (t *Tensor) Argsort(ctx ml.Context) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_argsort(ctx.(*Context).ctx, t.t, C.GGML_SORT_ORDER_ASC),
	}
}
Michael Yang's avatar
Michael Yang committed
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335

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

func (t *Tensor) Variance(ctx ml.Context) ml.Tensor {
	return t.Add(ctx, t.Mean(ctx).Scale(ctx, -1)).
		Sqr(ctx).
		SumRows(ctx).
		Scale(ctx, 1/float64(t.Dim(0)))
}

func (t *Tensor) Stddev(ctx ml.Context) ml.Tensor {
	return t.Variance(ctx).Sqrt(ctx)
}

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

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

func (t *Tensor) Clamp(ctx ml.Context, min, max float32) ml.Tensor {
	return &Tensor{
		b: t.b,
		t: C.ggml_clamp(ctx.(*Context).ctx, t.t, C.float(min), C.float(max)),
	}
}