makarna/pkg/compute/scratch_set.go

package compute

import "fmt"

type ScratchSet struct {
	spaces map[int]*ScratchSpace
}

func NewScratchSet(gpus []int, bytes int) (*ScratchSet, error) {
	ss := &ScratchSet{spaces: make(map[int]*ScratchSpace)}
	for _, gpu := range gpus {
		if gpu < 0 {
			continue
		}
		if _, ok := ss.spaces[gpu]; ok {
			continue
		}
		sc, err := NewScratchSpace(gpu, bytes)
		if err != nil {
			ss.Free()
			return nil, err
		}
		ss.spaces[gpu] = sc
	}
	if ss.spaces == nil {
		return nil, fmt.Errorf("scratch set init failed")
	}
	return ss, nil
}

func (s *ScratchSet) Scratch(gpu int) *ScratchSpace {
	if s == nil || s.spaces == nil {
		return nil
	}
	return s.spaces[gpu]
}

func (s *ScratchSet) Reset() {
	if s == nil || s.spaces == nil {
		return
	}
	for _, sc := range s.spaces {
		sc.Reset()
	}
}

func (s *ScratchSet) Free() {
	if s == nil || s.spaces == nil {
		return
	}
	for _, sc := range s.spaces {
		sc.Free()
	}
	s.spaces = nil
}