Add allocator abstraction

1 files changed, 400 insertions, 0 deletions

diff --git a/gpt2/allocator.go b/gpt2/allocator.go
new file mode 100644
index 0000000..e65c752
--- /dev/null
+++ b/gpt2/allocator.go
@@ -0,0 +1,400 @@
+package gpt2
+
+import (
+	"fmt"
+
+	ort "github.com/yalue/onnxruntime_go"
+)
+
+type Allocator struct {
+	step        int64
+	inputNames  []string
+	outputNames []string
+	values      map[string]ort.Value
+	withCache   bool
+}
+
+func NewAllocator(withCache bool) *Allocator {
+	return &Allocator{
+		values:    make(map[string]ort.Value),
+		withCache: withCache,
+	}
+}
+
+func (a *Allocator) InputNames() []string {
+	names := make([]string, 0, 3+2*nLayers)
+
+	names = append(names, "input_ids", "position_ids", "attention_mask")
+
+	if a.withCache {
+		for i := range nLayers {
+			names = append(names, fmt.Sprintf("past_key_values.%d.key", i), fmt.Sprintf("past_key_values.%d.value", i))
+		}
+	}
+
+	return names
+}
+
+func (a *Allocator) OutputNames() []string {
+	names := make([]string, 0, 1+2*nLayers)
+
+	names = append(names, "logits")
+
+	if a.withCache {
+		for i := range nLayers {
+			names = append(names, fmt.Sprintf("present.%d.key", i), fmt.Sprintf("present.%d.value", i))
+		}
+	}
+
+	return names
+}
+
+func (a *Allocator) Init(tokens []int64) error {
+	a.Destroy()
+
+	a.values = make(map[string]ort.Value)
+	a.step = 0
+
+	if err := a.initInputs(tokens); err != nil {
+		a.Destroy()
+
+		return err
+	}
+
+	if err := a.initOutputs(tokens); err != nil {
+		a.Destroy()
+
+		return err
+	}
+
+	a.step = int64(len(tokens))
+
+	return nil
+}
+
+func (a *Allocator) initInputs(tokens []int64) error {
+	capacity := 3
+
+	if a.withCache {
+		capacity += 2 * nLayers
+	}
+
+	names := make([]string, 0, capacity)
+
+	if err := a.inputIDs(tokens, false); err != nil {
+		return err
+	}
+
+	names = append(names, "input_ids")
+
+	if err := a.positionIDs(tokens, 0, false); err != nil {
+		return err
+	}
+
+	names = append(names, "position_ids")
+
+	if err := a.attentionMask(tokens, 0, false); err != nil {
+		return err
+	}
+
+	names = append(names, "attention_mask")
+
+	if !a.withCache {
+		a.inputNames = names
+
+		return nil
+	}
+
+	for i := range int64(nLayers) {
+		if err := a.pastKeyValues(i, "key", false); err != nil {
+			return err
+		}
+
+		names = append(names, fmt.Sprintf("past_key_values.%d.key", i))
+
+		if err := a.pastKeyValues(i, "value", false); err != nil {
+			return err
+		}
+
+		names = append(names, fmt.Sprintf("past_key_values.%d.value", i))
+	}
+
+	a.inputNames = names
+
+	return nil
+}
+
+func (a *Allocator) initOutputs(tokens []int64) error {
+	capacity := 1
+
+	if a.withCache {
+		capacity += 2 * nLayers
+	}
+
+	names := make([]string, 0, capacity)
+
+	if err := a.logits(tokens, false); err != nil {
+		return err
+	}
+
+	names = append(names, "logits")
+
+	if !a.withCache {
+		a.outputNames = names
+
+		return nil
+	}
+
+	for i := range int64(nLayers) {
+		if err := a.presentKeyValues(tokens, 0, i, "key", false); err != nil {
+			return err
+		}
+
+		names = append(names, fmt.Sprintf("present.%d.key", i))
+
+		if err := a.presentKeyValues(tokens, 0, i, "value", false); err != nil {
+			return err
+		}
+
+		names = append(names, fmt.Sprintf("present.%d.value", i))
+	}
+
+	a.outputNames = names
+
+	return nil
+}
+
+func (a *Allocator) Step(token int64) error {
+	tokens := []int64{token}
+
+	if err := a.inputIDs(tokens, true); err != nil {
+		return err
+	}
+
+	if err := a.positionIDs(tokens, a.step, true); err != nil {
+		return err
+	}
+
+	if err := a.attentionMask(tokens, a.step, true); err != nil {
+		return err
+	}
+
+	if err := a.logits(tokens, true); err != nil {
+		return err
+	}
+
+	for i := range int64(nLayers) {
+		for _, suffix := range []string{"key", "value"} {
+			if err := a.rotateCache(tokens, i, suffix); err != nil {
+				return err
+			}
+		}
+	}
+
+	a.step++
+
+	return nil
+}
+
+func (a *Allocator) Destroy() {
+	for _, v := range a.values {
+		_ = v.Destroy()
+	}
+}
+
+func (a *Allocator) Inputs() ([]string, []ort.Value) {
+	vals := make([]ort.Value, len(a.inputNames))
+
+	for i, n := range a.inputNames {
+		vals[i] = a.values[n]
+	}
+
+	return a.inputNames, vals
+}
+
+func (a *Allocator) Outputs() ([]string, []ort.Value) {
+	vals := make([]ort.Value, len(a.outputNames))
+
+	for i, n := range a.outputNames {
+		vals[i] = a.values[n]
+	}
+
+	return a.outputNames, vals
+}
+
+func (a *Allocator) inputIDs(tokens []int64, force bool) error {
+	const name = "input_ids"
+
+	if _, ok := a.values[name]; ok {
+		if !force {
+			panic("input_ids already allocated")
+		}
+
+		_ = a.values[name].Destroy()
+	}
+
+	shape := []int64{1, int64(len(tokens))}
+
+	if t, err := ort.NewTensor[int64](shape, tokens); err != nil {
+		return err
+	} else {
+		a.values[name] = ort.Value(t)
+	}
+
+	return nil
+}
+
+func (a *Allocator) positionIDs(tokens []int64, start int64, force bool) error {
+	const name = "position_ids"
+
+	if _, ok := a.values[name]; ok {
+		if !force {
+			panic("position_ids already allocated")
+		}
+
+		_ = a.values[name].Destroy()
+	}
+
+	data := make([]int64, len(tokens))
+	shape := []int64{1, int64(len(tokens))}
+
+	for i := range len(tokens) {
+		data[i] = start + int64(i)
+	}
+
+	if t, err := ort.NewTensor[int64](shape, data); err != nil {
+		return err
+	} else {
+		a.values[name] = ort.Value(t)
+	}
+
+	return nil
+}
+
+func (a *Allocator) attentionMask(tokens []int64, start int64, force bool) error {
+	const name = "attention_mask"
+
+	if _, ok := a.values[name]; ok {
+		if !force {
+			panic("attention_mask already allocated")
+		}
+
+		_ = a.values[name].Destroy()
+	}
+
+	data := make([]int64, start+int64(len(tokens)))
+	shape := []int64{1, start + int64(len(tokens))}
+
+	for i := range data {
+		data[i] = 1
+	}
+
+	if t, err := ort.NewTensor[int64](shape, data); err != nil {
+		return err
+	} else {
+		a.values[name] = ort.Value(t)
+	}
+
+	return nil
+}
+
+func (a *Allocator) pastKeyValues(i int64, suffix string, force bool) error {
+	if i > nLayers {
+		panic("invalid layer index")
+	}
+
+	if suffix != "key" && suffix != "value" {
+		panic("invalid suffix")
+	}
+
+	name := fmt.Sprintf("past_key_values.%d.%s", i, suffix)
+
+	if _, ok := a.values[name]; ok {
+		if !force {
+			panic(name + " already allocated")
+		}
+
+		_ = a.values[name].Destroy()
+	}
+
+	shape := []int64{1, int64(nHeads), 0, int64(headDim)}
+
+	if t, err := ort.NewEmptyTensor[float32](shape); err != nil {
+		return err
+	} else {
+		a.values[name] = ort.Value(t)
+	}
+
+	return nil
+}
+
+func (a *Allocator) logits(tokens []int64, force bool) error {
+	const name = "logits"
+
+	if _, ok := a.values[name]; ok {
+		if !force {
+			panic("logits already allocated")
+		}
+
+		_ = a.values[name].Destroy()
+	}
+
+	shape := []int64{1, int64(len(tokens)), int64(vocabSize)}
+
+	if t, err := ort.NewEmptyTensor[float32](shape); err != nil {
+		return err
+	} else {
+		a.values[name] = ort.Value(t)
+	}
+
+	return nil
+}
+
+func (a *Allocator) presentKeyValues(tokens []int64, start, i int64, suffix string, force bool) error {
+	if i > nLayers {
+		panic("invalid layer index")
+	}
+
+	if suffix != "key" && suffix != "value" {
+		panic("invalid suffix")
+	}
+
+	name := fmt.Sprintf("present.%d.%s", i, suffix)
+
+	if _, ok := a.values[name]; ok {
+		if !force {
+			panic(name + " already allocated")
+		}
+
+		_ = a.values[name].Destroy()
+	}
+
+	shape := []int64{1, int64(nHeads), start + int64(len(tokens)), int64(headDim)}
+
+	if t, err := ort.NewEmptyTensor[float32](shape); err != nil {
+		return err
+	} else {
+		a.values[name] = ort.Value(t)
+	}
+
+	return nil
+}
+
+func (a *Allocator) rotateCache(tokens []int64, i int64, suffix string) error {
+	past := fmt.Sprintf("past_key_values.%d.%s", i, suffix)
+	present := fmt.Sprintf("present.%d.%s", i, suffix)
+
+	if v, ok := a.values[past]; ok {
+		_ = v.Destroy()
+	}
+
+	a.values[past] = a.values[present]
+
+	delete(a.values, present)
+
+	if err := a.presentKeyValues(tokens, a.step, i, suffix, false); err != nil {
+		return err
+	}
+
+	return nil
+}