ggml.go 32.3 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
139
140
141
	requiredMemory.CPU.Name = C.GoString(C.ggml_backend_dev_name(cpuDeviceBufferType.d))
	requiredMemory.CPU.Weights = make([]ml.Memory, blocks+1)
	requiredMemory.CPU.Cache = make([]ml.Memory, blocks+1)

Michael Yang's avatar
Michael Yang committed
142
	// create list of buffer types for each gpu
143
	var gpuDeviceBufferTypes []deviceBufferType
144
145
	requiredMemory.GPUs = make([]ml.DeviceMemory, len(gpus))
	for i, d := range gpus {
146
		bt := C.ggml_backend_dev_buffer_type(d)
147
		gpuDeviceBufferTypes = append(gpuDeviceBufferTypes, deviceBufferType{
148
			d:   d,
Michael Yang's avatar
Michael Yang committed
149
			bts: append([]*C.struct_ggml_backend_buffer_type{bt}, cpuDeviceBufferType.bts...),
150
		})
151
152
153
154
		btDeviceMemory[bt] = &requiredMemory.GPUs[i]
		requiredMemory.GPUs[i].Name = C.GoString(C.ggml_backend_dev_name(d))
		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
155
156
	}

Michael Yang's avatar
Michael Yang committed
157
158
159
160
161
	useDefaultSplit := true
	for _, s := range params.TensorSplit {
		if s != 0 {
			useDefaultSplit = false
			break
162
		}
Michael Yang's avatar
Michael Yang committed
163
	}
164

Michael Yang's avatar
Michael Yang committed
165
166
167
168
	// calculate splits
	splits := make([]float32, len(gpus))
	if useDefaultSplit {
		// default: split on free memory
169
170
171
172
173
		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
174
175
	} else {
		splits = params.TensorSplit
176
177
178
	}

	var sum float32
Michael Yang's avatar
Michael Yang committed
179
	// cumulative sum of all splits
180
181
182
183
184
	for i := range splits {
		sum += splits[i]
		splits[i] = sum
	}

Michael Yang's avatar
Michael Yang committed
185
	// normalize splits
186
	for i := range splits {
187
		splits[i] /= sum
188
189
	}

Michael Yang's avatar
Michael Yang committed
190
	// inputs always use cpu
Michael Yang's avatar
Michael Yang committed
191
	input := cpuDeviceBufferType
192

Michael Yang's avatar
Michael Yang committed
193
194
195
	// 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
196
	assignLayer := func(i int) deviceBufferType {
Michael Yang's avatar
Michael Yang committed
197
		if i < gpuRangeStart || i >= gpuRangeStop {
Michael Yang's avatar
Michael Yang committed
198
			return cpuDeviceBufferType
199
		}
200

Michael Yang's avatar
Michael Yang committed
201
		index := slices.IndexFunc(splits, func(f float32) bool { return float32(i-gpuRangeStart)/float32(gpuRangeStop-gpuRangeStart) < f })
202
		if index < 0 || index >= len(gpuDeviceBufferTypes) {
Michael Yang's avatar
Michael Yang committed
203
			return cpuDeviceBufferType
204
205
206
		}

		return gpuDeviceBufferTypes[index]
207
208
	}

Michael Yang's avatar
Michael Yang committed
209
	// repeating layers are assigned based on their index in reverse order, e.g. i / (block_count + 1)
210
	layers := make([]deviceBufferType, blocks)
211
	for i := range layers {
212
		layers[i] = assignLayer(i)
213
214
	}

Michael Yang's avatar
Michael Yang committed
215
	// outputs are assigned iff allowed by splits and configured number of gpu layers
216
	output := assignLayer(blocks)
217
218
219

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

223
	type tensor struct {
224
		source *fsggml.Tensor
225
226
227
		target string
	}

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

Michael Yang's avatar
Michael Yang committed
231
	// contexts are shared by tensors of the same buffer type
232
	ctxs := make(map[*C.struct_ggml_backend_buffer_type]*C.struct_ggml_context)
233
	createTensor := func(t tensor, bts []*C.struct_ggml_backend_buffer_type, layer int) *C.struct_ggml_tensor {
234
235
236
237
238
239
240
		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
241

242
243
244
245
246
247
248
249
			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
250
			defer C.free(unsafe.Pointer(cname))
251
252
253
254
			if tt := C.ggml_get_tensor(ctxs[bt], cname); tt != nil {
				return tt
			}

255
			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
256
257
			C.ggml_set_name(tt, cname)

258
			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)))
259
260
261
262
263
264
265
266
267
268

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

269
270
271
272
273
			//nolint:staticcheck // TODO: check if buffer type supports this tensor
			return tt
		}

		return nil
Michael Yang's avatar
Michael Yang committed
274
275
	}

276
	contains := func(s string, parts ...string) bool {
277
278
279
280
281
282
283
284
		split := strings.Split(s, ".")
		for _, part := range parts {
			if slices.Contains(split, part) {
				return true
			}
		}

		return false
Michael Yang's avatar
Michael Yang committed
285
286
	}

287
288
	for _, t := range meta.Tensors().Items() {
		switch {
289
		case contains(t.Name, "position_embd", "token_embd", "token_norm_embd", "token_types"):
290
			createTensor(tensor{source: t}, input.bts, -1)
Michael Yang's avatar
Michael Yang committed
291
			if _, ok := meta.Tensors().GroupLayers()["output"]; !ok && t.Name == "token_embd.weight" {
292
				createTensor(tensor{source: t, target: "output.weight"}, output.bts, blocks)
Michael Yang's avatar
Michael Yang committed
293
			}
294
		case contains(t.Name, "cls", "output", "output_norm"):
295
			createTensor(tensor{source: t}, output.bts, blocks)
296
		case strings.HasPrefix(t.Name, "v.") || strings.HasPrefix(t.Name, "mm."):
Michael Yang's avatar
Michael Yang committed
297
			// TODO: assign vision tensors to the gpu if possible
298
			createTensor(tensor{source: t}, output.bts, blocks)
Michael Yang's avatar
Michael Yang committed
299
300
301
302
303
304
		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,
305
				}, layer.bts, i)
Michael Yang's avatar
Michael Yang committed
306
			}
307
		default:
Michael Yang's avatar
Michael Yang committed
308
309
310
311
			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
312
				}
Michael Yang's avatar
Michael Yang committed
313
			}
314

Michael Yang's avatar
Michael Yang committed
315
			if layerIndex >= 0 {
316
				createTensor(tensor{source: t}, layers[layerIndex].bts, layerIndex)
317
			} else {
Michael Yang's avatar
Michael Yang committed
318
				// load all other tensors on the cpu
319
				createTensor(tensor{source: t}, input.bts, -1)
320
321
322
			}
		}
	}
Michael Yang's avatar
Michael Yang committed
323

Michael Yang's avatar
Michael Yang committed
324
325
	// allocate buffers for each context
	bbs := make(map[*C.struct_ggml_context]*C.struct_ggml_backend_buffer, len(ctxs))
326
327
328
329
330
331
	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)
332
333
334
335
336
337
338
339
340
341
		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
				}
			}
		}

342
		if b == nil {
343
			panic(ml.ErrNoMem{BackendMemory: requiredMemory})
344
345
		}

346
		C.ggml_backend_buffer_set_usage(b, C.GGML_BACKEND_BUFFER_USAGE_WEIGHTS)
Michael Yang's avatar
Michael Yang committed
347
		bbs[c] = b
348
349
350
	}

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

Michael Yang's avatar
Michael Yang committed
354
	// map tensor names to tensors for easy lookup later
355
356
357
358
359
360
361
	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
		}
	}

362
363
364
365
366
367
368
369
370
371
372
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
	// 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
		}(),
408
409
410
		requiredMemory: &requiredMemory,
		btDeviceMemory: btDeviceMemory,
		maxGraphNodes:  maxGraphNodes,
411
412
413
414
415
416
417
418
	}, nil
}

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

func (b *Backend) Load(ctx context.Context, progress func(float32)) error {
419
	var doneBytes atomic.Uint64
420
	totalBytes := uint64(b.meta.Length) - b.meta.Tensors().Offset
421
422
423

	g, ctx := errgroup.WithContext(ctx)
	g.SetLimit(runtime.GOMAXPROCS(0))
424
	for _, t := range b.meta.Tensors().Items() {
425
		t := t
426
		g.Go(func() error {
427
			tts := make([]*C.struct_ggml_tensor, max(1, len(b.tensorLoadTargets[t.Name])))
428
			for i := range tts {
429
				target := b.tensorLoadTargets[t.Name][i]
430
431
432
				if target == "" {
					target = t.Name
				}
433

434
				tt, ok := b.tensors[target]
435
436
437
				if !ok {
					return fmt.Errorf("unassigned tensor: %s", t.Name)
				}
Michael Yang's avatar
Michael Yang committed
438

439
440
441
				tts[i] = tt
			}

442
443
			// 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.
444
			file, err := os.Open(b.modelPath)
445
			if err != nil {
446
				slog.Warn("file open error", "file", b.modelPath, "error", err)
447
448
449
				return err
			}
			defer file.Close()
450
			sr := io.NewSectionReader(file, int64(b.meta.Tensors().Offset+t.Offset), int64(t.Size()))
451
452
453
454
			bts := make([]byte, 128*format.KibiByte)

			var s uint64
			for s < t.Size() {
455
456
457
458
459
				// 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
				}

460
461
				n, err := io.ReadFull(sr, bts[:min(len(bts), int(t.Size()-s))])
				if err != nil {
462
					slog.Warn("file read error", "file", b.modelPath, "error", err)
463
					return err
464
				}
Michael Yang's avatar
Michael Yang committed
465

466
467
				for _, tt := range tts {
					C.ggml_backend_tensor_set(tt, unsafe.Pointer(&bts[0]), C.size_t(s), C.size_t(n))
468
				}
Michael Yang's avatar
Michael Yang committed
469

470
471
				s += uint64(n)

472
				if progress != nil {
473
					done := doneBytes.Add(uint64(n))
474
					progress(float32(done) / float32(totalBytes))
475
476
477
478
479
				}
			}

			return nil
		})
Michael Yang's avatar
Michael Yang committed
480
481
	}

482
	if err := g.Wait(); err != nil {
483
		return err
484
485
	}

486
	return nil
Michael Yang's avatar
Michael Yang committed
487
488
}

489
490
491
492
func (b *Backend) BackendMemory() ml.BackendMemory {
	return *b.requiredMemory
}

493
func (b *Backend) Config() fs.Config {
Michael Yang's avatar
Michael Yang committed
494
495
496
497
	return b.meta.KV()
}

func (b *Backend) Get(name string) ml.Tensor {
498
499
	if t, ok := b.tensors[name]; ok {
		return &Tensor{b: b, t: t}
Michael Yang's avatar
Michael Yang committed
500
501
502
503
504
505
	}

	return nil
}

func (b *Backend) NewContext() ml.Context {
Michael Yang's avatar
Michael Yang committed
506
	return b.NewContextSize(b.maxGraphNodes)
507
508
509
}

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

514
515
	var allocatedBuffers []*C.struct_ggml_backend_buffer

Michael Yang's avatar
Michael Yang committed
516
	return &Context{
517
518
		b:             b,
		maxGraphNodes: n,
519
		ctx: C.ggml_init(C.struct_ggml_init_params{
520
			mem_size: C.size_t(n)*C.ggml_tensor_overhead() + C.ggml_graph_overhead_custom(C.size_t(n), false),
521
522
			no_alloc: true,
		}),
523
		allocatedBuffers: &allocatedBuffers,
524
		layer:            -1,
Michael Yang's avatar
Michael Yang committed
525
526
527
	}
}

528
func (b *Backend) CacheConfig() ml.CacheConfig {
529
530
531
532
533
	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}
	}
534
535
}

Michael Yang's avatar
Michael Yang committed
536
type Context struct {
537
	b *Backend
Michael Yang's avatar
Michael Yang committed
538

539
	ctx   *C.struct_ggml_context
Michael Yang's avatar
Michael Yang committed
540
	graph *C.struct_ggml_cgraph
541

542
543
	// buft is the buffer type used for new tensors
	buft *C.struct_ggml_backend_buffer_type
544

545
546
547
548
	// 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
549
	// maxGraphNodes is the maximum allowed number of graph nodes in this context
550
	maxGraphNodes int
551
552
553

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

556
func (c *Context) Input() ml.Context {
Michael Yang's avatar
Michael Yang committed
557
	if c.b.input != nil {
558
		return &Context{
559
560
561
562
563
			b:                c.b,
			ctx:              c.ctx,
			buft:             c.b.input,
			allocatedBuffers: c.allocatedBuffers,
			maxGraphNodes:    c.maxGraphNodes,
564
			layer:            -1,
565
566
567
		}
	}

568
	return c
569
570
}

571
func (c *Context) Layer(i int) ml.Context {
572
	if buft, ok := c.b.layers[i]; ok {
573
		return &Context{
574
575
576
577
578
			b:                c.b,
			ctx:              c.ctx,
			buft:             buft,
			allocatedBuffers: c.allocatedBuffers,
			maxGraphNodes:    c.maxGraphNodes,
579
			layer:            i,
580
581
582
		}
	}

583
	return c
584
585
}

586
func (c *Context) Forward(tensors ...ml.Tensor) ml.Context {
Michael Yang's avatar
Michael Yang committed
587
	if c.graph == nil {
588
		c.graph = C.ggml_new_graph_custom(c.ctx, C.size_t(c.maxGraphNodes), false)
Michael Yang's avatar
Michael Yang committed
589
590
	}

591
592
593
594
595
	for _, tensor := range tensors {
		C.ggml_build_forward_expand(c.graph, tensor.(*Tensor).t)
	}

	return c
Michael Yang's avatar
Michael Yang committed
596
597
}

598
func (c *Context) Compute(tensors ...ml.Tensor) {
599
	C.ggml_backend_sched_graph_compute_async(c.b.sched, c.graph)
Michael Yang's avatar
Michael Yang committed
600
	C.ggml_backend_sched_reset(c.b.sched)
Michael Yang's avatar
Michael Yang committed
601

602
603
604
	needSync := true
	sync := func() {
		if needSync {
605
			C.ggml_backend_sched_synchronize(c.b.sched)
606
607
608
			needSync = false
		}
	}
Michael Yang's avatar
Michael Yang committed
609

610
611
612
	for _, t := range tensors {
		if C.ggml_nbytes(t.(*Tensor).t) > 0 {
			t.(*Tensor).sync = sync
613
614
		}
	}
Michael Yang's avatar
Michael Yang committed
615
616
}

617
618
func (c *Context) Reserve() {
	reserved := C.ggml_backend_sched_reserve(c.b.sched, c.graph)
619
620

	slog.Debug("compute graph", "nodes", C.ggml_graph_n_nodes(c.graph), "splits", C.ggml_backend_sched_get_n_splits(c.b.sched))
621
622
623
624
625
626

	// 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{}
	}

627
	for i := range c.b.schedBackends {
628
629
630
631
632
633
634
635
636
637
		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
		}

638
		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])),
639
			"size", format.HumanBytes2(uint64(bufferStatus.size)))
640
641
	}

642
643
644
	if !reserved {
		panic(ml.ErrNoMem{BackendMemory: *c.b.requiredMemory})
	}
645
646
}

647
func (c *Context) MaxGraphNodes() int {
648
	return c.maxGraphNodes
Jesse Gross's avatar
Jesse Gross committed
649
650
}

651
652
653
func shapeToGGML(shape []int) *C.int64_t {
	sh := make([]C.int64_t, len(shape))
	for i, s := range shape {
654
		sh[i] = C.int64_t(s)
655
656
657
658
659
	}

	return &sh[0]
}

660
661
662
663
func pad(length, pad C.size_t) C.size_t {
	return ((length + pad - 1) / pad) * pad
}

664
func (c *Context) newTensor(dtype ml.DType, shape []int) ml.Tensor {
665
	if c.buft == nil {
666
		panic("set Input or Layer before creating tensors")
667
668
	}

Michael Yang's avatar
Michael Yang committed
669
670
671
672
673
674
	var cdtype uint32
	switch dtype {
	case ml.DTypeF32:
		cdtype = C.GGML_TYPE_F32
	case ml.DTypeF16:
		cdtype = C.GGML_TYPE_F16
675
676
677
678
	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
679
680
681
682
683
684
	case ml.DTypeI32:
		cdtype = C.GGML_TYPE_I32
	default:
		panic("unsupported dtype")
	}

Jesse Gross's avatar
Jesse Gross committed
685
	if len(shape) < 1 || shape[0] == 0 {
Michael Yang's avatar
Michael Yang committed
686
		var shape C.int64_t = 0
687
		return &Tensor{b: c.b, t: C.ggml_new_tensor(c.ctx, cdtype, 1, &shape)}
Michael Yang's avatar
Michael Yang committed
688
	} else if len(shape) > 4 {
Michael Yang's avatar
Michael Yang committed
689
690
691
692
693
694
695
696
697
		panic("unsupported number of dimensions")
	}

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

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

701
	b := C.ggml_backend_buft_alloc_buffer(c.buft, size)
702
703
704
705
706
707
708
709
710
711
712
	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
		}
	}

713
	if b == nil {
714
		panic(ml.ErrNoMem{BackendMemory: *c.b.requiredMemory})
715
716
	}

717
	*c.allocatedBuffers = append(*c.allocatedBuffers, b)
Michael Yang's avatar
Michael Yang committed
718
	C.ggml_backend_tensor_alloc(b, t, C.ggml_backend_buffer_get_base(b))
719
	return &Tensor{b: c.b, t: t}
720
721
}

722
func (c *Context) Empty(dtype ml.DType, shape ...int) ml.Tensor {
723
	return c.newTensor(dtype, shape)
724
725
}

726
func (c *Context) Zeros(dtype ml.DType, shape ...int) ml.Tensor {
727
	t := c.newTensor(dtype, shape)
728
729
	C.ggml_set_zero(t.(*Tensor).t)
	return t
Michael Yang's avatar
Michael Yang committed
730
731
}

732
func checkShape[S ~[]E, E any](s S, shape ...int) {
Michael Yang's avatar
Michael Yang committed
733
	n := len(s)
Jesse Gross's avatar
Jesse Gross committed
734
735

	if n == 0 {
736
		return
Jesse Gross's avatar
Jesse Gross committed
737
738
	}

Michael Yang's avatar
Michael Yang committed
739
740
741
742
743
	for _, v := range shape {
		n /= v
	}

	if n != 1 {
744
		panic(fmt.Errorf("invalid shape: %v", shape))
Michael Yang's avatar
Michael Yang committed
745
746
747
	}
}

748
749
func (c *Context) FromFloatSlice(s []float32, shape ...int) ml.Tensor {
	checkShape(s, shape...)
750

751
	t := c.newTensor(ml.DTypeF32, shape)
752

Jesse Gross's avatar
Jesse Gross committed
753
754
755
756
	if len(s) > 0 {
		C.ggml_backend_tensor_set(t.(*Tensor).t, unsafe.Pointer(&s[0]), 0, C.ggml_nbytes(t.(*Tensor).t))
	}

757
	return t
Michael Yang's avatar
Michael Yang committed
758
759
}

760
761
func (c *Context) FromIntSlice(s []int32, shape ...int) ml.Tensor {
	checkShape(s, shape...)
762

763
	t := c.newTensor(ml.DTypeI32, shape)
764

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

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

Michael Yang's avatar
arange  
Michael Yang committed
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
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))
		}

787
		return c.Input().FromIntSlice(arange, len(arange))
Michael Yang's avatar
arange  
Michael Yang committed
788
789
790
791
792
	default:
		panic("unsupported dtype for arange")
	}
}

Michael Yang's avatar
Michael Yang committed
793
794
func (c *Context) Close() {
	if c != nil {
795
796
797
798
799
		for _, b := range *c.allocatedBuffers {
			C.ggml_backend_buffer_free(b)
		}
		*c.allocatedBuffers = nil

800
801
		C.ggml_free(c.ctx)
	}
Michael Yang's avatar
Michael Yang committed
802
803
804
}

type Tensor struct {
805
	b    *Backend
Michael Yang's avatar
Michael Yang committed
806
	t    *C.struct_ggml_tensor
807
	sync func()
Michael Yang's avatar
Michael Yang committed
808
809
810
811
812
813
814
815
816
817
}

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

818
819
func (t *Tensor) Dim(n int) int {
	return int(t.t.ne[n])
Michael Yang's avatar
Michael Yang committed
820
821
}

822
823
func (t *Tensor) Stride(n int) int {
	return int(t.t.nb[n])
Michael Yang's avatar
Michael Yang committed
824
825
}

826
827
func (t *Tensor) Shape() []int {
	shape := make([]int, C.ggml_n_dims(t.t))
Michael Yang's avatar
Michael Yang committed
828
829
830
831
832
833
834
	for i := range shape {
		shape[i] = t.Dim(i)
	}

	return shape
}

835
836
837
838
839
840
841
842
843
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
844
845
}

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

	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
861
862
	case C.GGML_TYPE_F16:
		return ml.DTypeF16
863
864
865
866
	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
867
868
869
870
871
872
873
	case C.GGML_TYPE_I32:
		return ml.DTypeI32
	default:
		return ml.DTypeOther
	}
}

874
875
876
877
878
879
880
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
881
882
func (t *Tensor) Add(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	return &Tensor{
883
		b: t.b,
Michael Yang's avatar
Michael Yang committed
884
885
886
887
		t: C.ggml_add(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
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
909
910
911
912
913
914
915
916
917
918
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{
919
		b: t.b,
Michael Yang's avatar
Michael Yang committed
920
921
922
923
924
925
		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{
926
		b: t.b,
Michael Yang's avatar
Michael Yang committed
927
928
929
930
931
932
		t: C.ggml_cont(ctx.(*Context).ctx, t.t),
	}
}

func (t *Tensor) Mul(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	return &Tensor{
933
		b: t.b,
Michael Yang's avatar
Michael Yang committed
934
935
936
937
		t: C.ggml_mul(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

938
939
940
941
942
943
944
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
945
946
func (t *Tensor) Mulmat(ctx ml.Context, t2 ml.Tensor) ml.Tensor {
	return &Tensor{
947
		b: t.b,
Michael Yang's avatar
Michael Yang committed
948
949
950
951
		t: C.ggml_mul_mat(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

952
953
954
955
956
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{
957
		b: t.b,
958
959
960
961
		t: mul,
	}
}

Michael Yang's avatar
llama4  
Michael Yang committed
962
963
964
965
966
967
968
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
969
func (t *Tensor) LayerNorm(ctx ml.Context, w, b ml.Tensor, eps float32) ml.Tensor {
Michael Yang's avatar
llama4  
Michael Yang committed
970
971
972
973
974
975
	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
976
977
	}

Michael Yang's avatar
llama4  
Michael Yang committed
978
	return &Tensor{b: t.b, t: tt}
Michael Yang's avatar
Michael Yang committed
979
980
981
}

func (t *Tensor) RMSNorm(ctx ml.Context, w ml.Tensor, eps float32) ml.Tensor {
Michael Yang's avatar
llama4  
Michael Yang committed
982
983
984
985
986
987
	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
988
989
}

990
func (t *Tensor) Pad(ctx ml.Context, shape ...int) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
991
992
	if len(shape) != 4 {
		panic("expected 4 dimensions")
993
994
	} else if shape[3] != 0 {
		panic("cuda does not support 4d tensors")
Michael Yang's avatar
Michael Yang committed
995
996
997
	}

	return &Tensor{
998
		b: t.b,
Michael Yang's avatar
Michael Yang committed
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
		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{
1009
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1010
1011
1012
1013
1014
1015
		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{
1016
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1017
1018
1019
1020
1021
1022
		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{
1023
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1024
1025
1026
1027
		t: C.ggml_cpy(ctx.(*Context).ctx, t.t, t2.(*Tensor).t),
	}
}

1028
func (t *Tensor) Reshape(ctx ml.Context, shape ...int) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
1029
1030
1031
	switch len(shape) {
	case 1:
		return &Tensor{
1032
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1033
1034
1035
1036
			t: C.ggml_reshape_1d(ctx.(*Context).ctx, t.t, C.int64_t(shape[0])),
		}
	case 2:
		return &Tensor{
1037
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1038
1039
1040
1041
			t: C.ggml_reshape_2d(ctx.(*Context).ctx, t.t, C.int64_t(shape[0]), C.int64_t(shape[1])),
		}
	case 3:
		return &Tensor{
1042
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1043
1044
1045
1046
			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{
1047
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1048
1049
1050
1051
1052
1053
1054
1055
1056
			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{
1057
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1058
1059
1060
1061
		t: C.ggml_scale(ctx.(*Context).ctx, t.t, (C.float)(s)),
	}
}

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

1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
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
1090
1091
func (t *Tensor) Tanh(ctx ml.Context) ml.Tensor {
	return &Tensor{
1092
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1093
1094
1095
1096
		t: C.ggml_tanh_inplace(ctx.(*Context).ctx, t.t),
	}
}

Michael Yang's avatar
llama4  
Michael Yang committed
1097
1098
1099
1100
1101
1102
1103
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
1104
1105
1106
1107
func (t *Tensor) View(ctx ml.Context, offset int, shape ...int) ml.Tensor {
	switch len(shape) {
	case 1:
		return &Tensor{
1108
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1109
1110
1111
1112
			t: C.ggml_view_1d(ctx.(*Context).ctx, t.t, C.int64_t(shape[0]), C.size_t(offset)),
		}
	case 3:
		return &Tensor{
1113
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1114
1115
1116
1117
1118
1119
1120
			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{
1121
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1122
1123
1124
1125
1126
1127
1128
			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{
1129
			b: t.b,
Michael Yang's avatar
Michael Yang committed
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
			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")
	}
}

1140
func (t *Tensor) RoPE(ctx ml.Context, positions ml.Tensor, ropeDim int, ropeBase, ropeScale float32, options ...func(*rope.Options)) ml.Tensor {
1141
	// Default options
1142
	opts := &rope.Options{OriginalContextLength: 131072, Factors: &Tensor{}}
1143
1144
1145
1146
1147
1148

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

Jesse Gross's avatar
Jesse Gross committed
1149
1150
1151
1152
1153
	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
1154
	return &Tensor{
1155
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1156
		t: C.ggml_rope_ext(
1157
1158
			ctx.(*Context).ctx,
			dequant,
1159
1160
			positions.(*Tensor).t,
			opts.Factors.(*Tensor).t,
Michael Yang's avatar
Michael Yang committed
1161
			C.int(ropeDim),
1162
1163
			C.int(opts.Type),
			C.int(opts.OriginalContextLength),
Michael Yang's avatar
Michael Yang committed
1164
1165
			C.float(ropeBase),
			C.float(ropeScale),
1166
1167
1168
1169
			C.float(0.0),
			C.float(1.0),
			C.float(32.0),
			C.float(1.0),
Michael Yang's avatar
Michael Yang committed
1170
1171
1172
1173
		),
	}
}

1174
1175
1176
1177
1178
1179
1180
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
1181
1182
func (t *Tensor) GELU(ctx ml.Context) ml.Tensor {
	return &Tensor{
1183
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1184
1185
1186
1187
1188
1189
		t: C.ggml_gelu_inplace(ctx.(*Context).ctx, t.t),
	}
}

func (t *Tensor) SILU(ctx ml.Context) ml.Tensor {
	return &Tensor{
1190
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1191
1192
1193
1194
1195
1196
		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{
1197
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1198
1199
1200
		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)),
	}
}
1201

Michael Yang's avatar
Michael Yang committed
1202
func (t *Tensor) AvgPool2D(ctx ml.Context, k, s int, p float32) ml.Tensor {
Michael Yang's avatar
Michael Yang committed
1203
1204
	return &Tensor{
		b: t.b,
Michael Yang's avatar
Michael Yang committed
1205
		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
1206
1207
1208
	}
}

Michael Yang's avatar
Michael Yang committed
1209
1210
1211
1212
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
1213
		tt = C.ggml_set_1d(ctx.(*Context).ctx, t.t, t2.(*Tensor).t, C.size_t(offset))
Michael Yang's avatar
Michael Yang committed
1214
	case 1:
Michael Yang's avatar
Michael Yang committed
1215
		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
1216
1217
1218
1219
1220
1221
1222
	default:
		panic("unsupported number of dimensions")
	}

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

1223
1224
1225
1226
1227
1228
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
	}

1229
1230
1231
	query := t.Permute(ctx, 0, 2, 1, 3)
	key = key.Permute(ctx, 0, 2, 1, 3)

1232
1233
	if t.b.flashAttention {
		value = value.Permute(ctx, 0, 2, 1, 3)
1234

1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
		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)
	}
1248
}
1249
1250
1251
1252
1253
1254
1255

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
1256
1257
1258
1259
1260
1261
1262

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)),
	}
}
1263
1264
1265
1266
1267
1268
1269

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