dtype.go

/*
 *	Copyright 2023 Jan Pfeifer
 *
 *	Licensed under the Apache License, Version 2.0 (the "License");
 *	you may not use this file except in compliance with the License.
 *	You may obtain a copy of the License at
 *
 *	http://www.apache.org/licenses/LICENSE-2.0
 *
 *	Unless required by applicable law or agreed to in writing, software
 *	distributed under the License is distributed on an "AS IS" BASIS,
 *	WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *	See the License for the specific language governing permissions and
 *	limitations under the License.
 */

package shapes

import (
	"github.com/gomlx/gomlx/types/exceptions"
	"github.com/pkg/errors"
	"math"
	"reflect"
	"strconv"
	"unsafe"
)

// DType indicates the type of the unit element of a Tensor (or its representation in
// a computation graph). It enumerates the known data types. So far only
// Bool, Uint8 (U8), Int32 (I32), Int64 (I64), Uint64 (U64), Float32 (F32) and Float64 (F64) are supported.
//
// The values of DType must match "tensorflow/compiler/xla/xla_data.pb.h", hence it needs
// to be an int32.
//
// See example in package shapes documentation.
type DType int32

//go:generate enumer -type=DType -yaml -json -text -values

// DType constants must match `tensorflow/compiler/xla/xla_data.proto`.
const (
	InvalidDType DType = iota
	Bool               // Bool, but also known as PRED in `xla_data.proto`.
	Int8               // S8
	Int16              // S16
	Int32              // S32
	Int64              // S64
	UInt8              // U8
	UInt16             // U16
	UInt32             // U32
	UInt64             // U64
	Float16            // F16
	Float32            // F32
	Float64            // F64

	BFloat16   DType = 16 // BF16
	Complex64  DType = 15 // C64
	Complex128 DType = 18 // C128

	Tuple      DType = 13
	OpaqueType DType = 14
	Token      DType = 17
)

// PRED type is an alias to Bool, used in `tensorflow/compiler/xla/xla_data.proto`.
const PRED = Bool

// Aliases:

const (
	U8   = UInt8
	U32  = UInt32
	U64  = UInt64
	I32  = Int32
	I64  = Int64
	F16  = Float16
	F32  = Float32
	F64  = Float64
	C64  = Complex64
	C128 = Complex128
)

// IsFloat returns whether dtype is a supported float -- float types not yet supported will return false.
// It returns false for complex numbers.
func (dtype DType) IsFloat() bool {
	return dtype == Float32 || dtype == Float64
}

// IsComplex returns whether dtype is a supported complex number type.
func (dtype DType) IsComplex() bool {
	return dtype == Complex64 || dtype == Complex128
}

// RealDType returns the real component of complex dtypes.
// For float dtypes, it returns itself.
//
// It returns InvalidDType for other non-(complex or float) dtypes.
func (dtype DType) RealDType() DType {
	if dtype.IsFloat() {
		return dtype
	}
	switch dtype {
	case Complex64:
		return Float32
	case Complex128:
		return Float64
	}
	// RealDType is not defined for other dtypes.
	return InvalidDType
}

// IsInt returns whether dtype is a supported integer type -- float types not yet supported will return false.
func (dtype DType) IsInt() bool {
	return dtype == Int64 || dtype == Int32 || dtype == UInt8 || dtype == UInt32 || dtype == UInt64
}

func (dtype DType) IsSupported() bool {
	return dtype == Bool || dtype == Float32 || dtype == Float64 ||
		dtype == Int64 || dtype == Int32 || dtype == UInt8 || dtype == UInt32 || dtype == UInt64 ||
		dtype == Complex64 || dtype == Complex128
}

// GoStr converts dtype to the corresponding Go type and convert that to string.
// Notice the names are different from the Dtype (so `Int64` dtype is simply `int` in Go).
func (dtype DType) GoStr() string {
	t := TypeForDType(dtype)
	return t.Name()
}

// Generated by `cmd/constraints_generator`:

// Supported lists the Go types that are supported by the graph package. Used as a Generics constraint.
// See also Number.
//
// Notice Go's `int` type is not portable, since it may translate to dtypes Int32 or Int64 depending
// on the platform.
//
// Generated by `cmd/constraints_generator`.
type Supported interface {
	bool | float32 | float64 | int | int32 | int64 | uint8 | uint32 | uint64 | complex64 | complex128
}

// Number represents the Go numeric types that are supported by graph package.
// Used as a Generics constraint.
// Notice that "int" becomes int64 in the implementation.
// Since it needs a 1:1 mapping, it doesn't support the native (Go) int64 type.
// It includes complex numbers.
type Number interface {
	float32 | float64 | int | int32 | int64 | uint8 | uint32 | uint64 | complex64 | complex128
}

// NumberNotComplex represents the Go numeric types that are supported by graph package except the complex numbers.
// Used as a Generics constraint.
// See Number for details.
type NumberNotComplex interface {
	float32 | float64 | int | int32 | int64 | uint8 | uint32 | uint64
}

// GoFloat represent a continuous Go numeric type, supported by GoMLX.
type GoFloat interface {
	float32 | float64
}

// Generated by `cmd/constraints_generator`:

// MultiDimensionSlice lists the Go types a Tensor can be converted to/from. There are no recursions in
// generics' constraint definitions, so we enumerate up to 7 levels of slices. Feel free to add
// more if needed, the implementation will work with any arbitrary number.
type MultiDimensionSlice interface {
	bool | float32 | float64 | int | int32 | int64 | uint8 | uint32 | uint64 | complex64 | complex128 |
		[]bool | []float32 | []float64 | []int | []int32 | []int64 | []uint8 | []uint32 | []uint64 | []complex64 | []complex128 |
		[][]bool | [][]float32 | [][]float64 | [][]int | [][]int32 | [][]int64 | [][]uint8 | [][]uint32 | [][]uint64 | [][]complex64 | [][]complex128 |
		[][][]bool | [][][]float32 | [][][]float64 | [][][]int | [][][]int32 | [][][]int64 | [][][]uint8 | [][][]uint32 | [][][]uint64 | [][][]complex64 | [][][]complex128 |
		[][][][]bool | [][][][]float32 | [][][][]float64 | [][][][]int | [][][][]int32 | [][][][]int64 | [][][][]uint8 | [][][][]uint32 | [][][][]uint64 | [][][][]complex64 | [][][][]complex128 |
		[][][][][]bool | [][][][][]float32 | [][][][][]float64 | [][][][][]int | [][][][][]int32 | [][][][][]int64 | [][][][][]uint8 | [][][][][]uint32 | [][][][][]uint64 | [][][][][]complex64 | [][][][][]complex128 |
		[][][][][][]bool | [][][][][][]float32 | [][][][][][]float64 | [][][][][][]int | [][][][][][]int32 | [][][][][][]int64 | [][][][][][]uint8 | [][][][][][]uint32 | [][][][][][]uint64 | [][][][][][]complex64 | [][][][][][]complex128
}

func DTypeGeneric[T Supported]() DType {
	var t T
	switch (any(t)).(type) {
	case float32:
		return Float32
	case float64:
		return Float64
	case int:
		switch strconv.IntSize {
		case 32:
			return Int32
		case 64:
			return Int64
		default:
			exceptions.Panicf("Cannot use int of %d bits with GoMLX -- try using int32 or int64", strconv.IntSize)
		}
	case int64:
		return Int64
	case int32:
		return Int32
	case bool:
		return Bool
	case uint8:
		return UInt8
	case uint32:
		return UInt32
	case uint64:
		return UInt64
	case complex64:
		return Complex64
	case complex128:
		return Complex128
	}
	return InvalidDType
}

// ConvertTo converts any scalar (typically returned by `tensor.Local.Value()`) of the
// supported dtypes to `T`.
// Returns 0 if value is not a scalar or not a supported number (e.g: bool).
// It doesn't work for if T (the output type) is a complex number.
// If value is a complex number, it converts by taking the real part of the number and
// discarding the imaginary part.
func ConvertTo[T NumberNotComplex](value any) T {
	switch v := value.(type) {
	case float32:
		return T(v)
	case float64:
		return T(v)
	case int:
		return T(v)
	case int64:
		return T(v)
	case int32:
		return T(v)
	case uint8:
		return T(v)
	case uint32:
		return T(v)
	case uint64:
		return T(v)
	case complex64:
		return T(real(v))
	case complex128:
		return T(real(v))
	}
	return T(0)
}

func DTypeForType(t reflect.Type) DType {
	switch t.Kind() {
	case reflect.Int:
		switch strconv.IntSize {
		case 32:
			return Int32
		case 64:
			return Int64
		default:
			exceptions.Panicf("cannot use int of %d bits with GoMLX -- try using int32 or int64", strconv.IntSize)
		}
	case reflect.Int64:
		return Int64
	case reflect.Int32:
		return Int32
	case reflect.Float32:
		return Float32
	case reflect.Float64:
		return Float64
	case reflect.Bool:
		return Bool
	case reflect.Uint8:
		return UInt8
	case reflect.Uint32:
		return UInt32
	case reflect.Uint64:
		return UInt64
	case reflect.Complex64:
		return Complex64
	case reflect.Complex128:
		return Complex128
	}
	return InvalidDType
}

// UnsafeSliceForDType creates a slice of the corresponding dtype
// and casts it to any.
// It uses unsafe.Slice.
// Set `len` to the number of `DType` elements (not the number of bytes).
func UnsafeSliceForDType(dtype DType, unsafePtr unsafe.Pointer, len int) any {
	switch dtype {
	case Int64:
		return unsafe.Slice((*int64)(unsafePtr), len)
	case Int32:
		return unsafe.Slice((*int32)(unsafePtr), len)
	case Float32:
		return unsafe.Slice((*float32)(unsafePtr), len)
	case Float64:
		return unsafe.Slice((*float64)(unsafePtr), len)
	case Bool:
		return unsafe.Slice((*bool)(unsafePtr), len)
	case UInt8:
		return unsafe.Slice((*uint8)(unsafePtr), len)
	case UInt32:
		return unsafe.Slice((*uint32)(unsafePtr), len)
	case UInt64:
		return unsafe.Slice((*uint64)(unsafePtr), len)
	case Complex64:
		return unsafe.Slice((*complex64)(unsafePtr), len)
	case Complex128:
		return unsafe.Slice((*complex128)(unsafePtr), len)
	}
	return nil
}

// TypeForDType returns the Go `reflect.Type` corresponding to the tensor DType.
func TypeForDType(dtype DType) reflect.Type {
	switch dtype {
	case Int64:
		return reflect.TypeOf(int64(0))
	case Int32:
		return reflect.TypeOf(int32(0))
	case Float32:
		return float32Type
	case Float64:
		return float64Type
	case Bool:
		return reflect.TypeOf(true)
	case UInt8:
		return reflect.TypeOf(uint8(0))
	case UInt32:
		return reflect.TypeOf(uint32(0))
	case UInt64:
		return reflect.TypeOf(uint64(0))
	case Complex64:
		return reflect.TypeOf(complex64(0))
	case Complex128:
		return reflect.TypeOf(complex128(0))
	}
	return reflect.TypeOf(nil)
}

// Type returns the Go `reflect.Type` corresponding to the tensor DType.
func (dtype DType) Type() reflect.Type {
	return TypeForDType(dtype)
}

// Memory returns the number of bytes in Go used to store the given DType.
func (dtype DType) Memory() int64 {
	t := TypeForDType(dtype)
	return int64(t.Size())
}

// Pre-generate constant reflect.TypeOf for convenience.
var (
	float32Type = reflect.TypeOf(float32(0))
	float64Type = reflect.TypeOf(float64(0))
)

// CastAsDType casts a numeric value to the corresponding for the DType.
// If the value is a slice it will convert to a newly allocated slice of
// the given DType.
//
// It doesn't work for complex numbers.
func CastAsDType(value any, dtype DType) any {
	typeOf := reflect.TypeOf(value)
	valueOf := reflect.ValueOf(value)
	newTypeOf := typeForSliceDType(typeOf, dtype)
	if typeOf.Kind() != reflect.Slice && typeOf.Kind() != reflect.Array {
		// Scalar value.
		if dtype == Bool {
			return !valueOf.IsZero()
		}
		if dtype == Complex64 {
			r := valueOf.Convert(float32Type).Interface().(float32)
			return complex(r, float32(0))
		}
		if dtype == Complex128 {
			r := valueOf.Convert(float64Type).Interface().(float64)
			return complex(r, float64(0))
		}
		// TODO: if adding support for non-native Go types (e.g: B16), we need
		//       to write our own conversion here.
		return valueOf.Convert(newTypeOf).Interface()
	}

	newValueOf := reflect.MakeSlice(newTypeOf, valueOf.Len(), valueOf.Len())
	for ii := 0; ii < valueOf.Len(); ii++ {
		elem := CastAsDType(valueOf.Index(ii).Interface(), dtype)
		newValueOf.Index(ii).Set(reflect.ValueOf(elem))
	}
	return newValueOf.Interface()
}

// typeForSliceDType recursively converts a type that is a (multi-dimension-) slice
// of some type, to a `reflect.Type` of a (multi-dimension-) slice of `dtype`.
// Arrays are converted to slices.
func typeForSliceDType(valueType reflect.Type, dtype DType) reflect.Type {
	if valueType.Kind() != reflect.Slice && valueType.Kind() != reflect.Array {
		// Base case for recursion, simply return the `reflect.Type` for the DType.
		return TypeForDType(dtype)
	}
	subType := typeForSliceDType(valueType.Elem(), dtype)
	return reflect.SliceOf(subType) // Return a slice of the recursively converted type.
}

// LowestValueForDType converted to the corresponding Go type.
// For float values it will return negative infinite.
// There is no lowest value for complex numbers, since they are not ordered.
func LowestValueForDType(dtype DType) any {
	switch dtype {
	case Int64:
		return int64(math.MinInt64)
	case Int32:
		return int32(math.MinInt32)
	case Float32:
		return float32(math.Inf(-1))
	case Float64:
		return math.Inf(-1)
	case Bool:
		return false
	case UInt8:
		return uint8(0)
	case UInt32:
		return uint32(0)
	case UInt64:
		return uint64(0)
	default:
		exceptions.Panicf("LowestValueForDType not defined for dtype %s", dtype)
	}
	return 0 // Never reaches here.
}

// SmallestNonZeroValueForDType is the smallest non-zero value dtypes.
// Only useful for float types.
// The return value is converted to the corresponding Go type.
// There is no smallest non-zero value for complex numbers, since they are not ordered.
func SmallestNonZeroValueForDType(dtype DType) any {
	switch dtype {
	case Int64:
		return int64(1)
	case Int32:
		return int32(1)
	case Float32:
		return float32(math.SmallestNonzeroFloat32)
	case Float64:
		return math.SmallestNonzeroFloat64
	case Bool:
		return true
	case UInt8:
		return uint8(1)
	case UInt32:
		return uint32(1)
	case UInt64:
		return uint64(1)
	}
	panic(errors.Errorf("SmallestNonZeroValueForDType not defined for dtype %s", dtype))
}