compute/optype.go at main · gomlx/compute · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0

package compute

// OpType is an enum of all generic operations that can be supported by a Backend.Builder.
//
// Notice: nothing precludes a specialized backend Builder to support other ops not included here.
// It requires some careful casting of interfaces by the caller.
type OpType int

//go:generate go tool enumer -type=OpType -trimprefix=OpType -output=gen_optype_enumer.go optype.go

const (
	OpTypeInvalid OpType = iota
	OpTypeParameter
	OpTypeConstant
	OpTypeIdentity
	OpTypeReduceWindow
	OpTypeRNGBitGenerator
	OpTypeBatchNormForInference
	OpTypeBatchNormForTraining
	OpTypeBatchNormGradient
	OpTypeBitCount

	OpTypeAbs
	OpTypeAdd
	OpTypeArgMinMax
	OpTypeAtan2
	OpTypeBitcast
	OpTypeBitwiseAnd
	OpTypeBitwiseNot
	OpTypeBitwiseOr
	OpTypeBitwiseXor
	OpTypeBroadcastInDim
	OpTypeCall
	OpTypeClamp
	OpTypeCeil
	OpTypeClz
	OpTypeComplex
	OpTypeConcatenate
	OpTypeConj
	OpTypeConvGeneral
	OpTypeConvertDType
	OpTypeCos
	OpTypeDiv
	OpTypeDot
	OpTypeDotGeneral
	OpTypeDynamicSlice
	OpTypeDynamicUpdateSlice
	OpTypeDynamicDimensionSize
	OpTypeDynamicShape
	OpTypeEqual
	OpTypeEqualTotalOrder
	OpTypeErf
	OpTypeExp
	OpTypeExpm1
	OpTypeFFT
	OpTypeFloor
	OpTypeGather
	OpTypeGreaterOrEqual
	OpTypeGreaterOrEqualTotalOrder
	OpTypeGreaterThan
	OpTypeGreaterThanTotalOrder
	OpTypeImag
	OpTypeIota
	OpTypeIsFinite
	OpTypeIsNaN
	OpTypeLessOrEqual
	OpTypeLessOrEqualTotalOrder
	OpTypeLessThan
	OpTypeLessThanTotalOrder
	OpTypeLog
	OpTypeLog1p
	OpTypeLogicalAnd
	OpTypeLogicalNot
	OpTypeLogicalOr
	OpTypeLogicalXor
	OpTypeLogistic
	OpTypeMax
	OpTypeMin
	OpTypeMul
	OpTypeNeg
	OpTypeNotEqual
	OpTypeNotEqualTotalOrder
	OpTypePad
	OpTypePow
	OpTypeReal
	OpTypeReduceBitwiseAnd
	OpTypeReduceBitwiseOr
	OpTypeReduceBitwiseXor
	OpTypeReduceLogicalAnd
	OpTypeReduceLogicalOr
	OpTypeReduceLogicalXor
	OpTypeReduceMax
	OpTypeReduceMin
	OpTypeReduceProduct
	OpTypeReduceSum
	OpTypeRem
	OpTypeReshape
	OpTypeReverse
	OpTypeRound
	OpTypeRsqrt
	OpTypeScatterMax
	OpTypeScatterMin
	OpTypeScatterSum
	OpTypeSelectAndScatterMax
	OpTypeSelectAndScatterMin
	OpTypeSelectAndScatterSum
	OpTypeShiftLeft
	OpTypeShiftRightArithmetic
	OpTypeShiftRightLogical
	OpTypeSign
	OpTypeSin
	OpTypeSlice
	OpTypeSqrt
	OpTypeSub
	OpTypeTanh
	OpTypeTranspose
	OpTypeWhere
	OpTypeOptimizationBarrier
	OpTypeSchedulingBarrier

	// Control flow operations

	OpTypeSort
	OpTypeWhile
	OpTypeIf

	// OpTypeCapturedValue represents a value captured from a parent scope in a closure.
	// This allows closures to reference values computed in enclosing functions.
	OpTypeCapturedValue

	// Collective (distributed across devices) operations

	OpTypeAllReduce
	OpTypeCollectiveBroadcast
	OpTypeAllGather

	// Internal operations (backend-specific, not part of public API)

	// OpTypeBlockForDotGeneral pre-blocks a tensor for efficient DotGeneral execution.
	// This is an internal optimization used by the go backend.
	OpTypeBlockForDotGeneral

	// Fused operations: high-level ops that backends may implement natively.
	// If supported (declared in Capabilities.Operations), GoMLX uses the
	// native implementation; otherwise it decomposes into primitives.

	OpTypeFusedSoftmax
	OpTypeFusedLayerNorm
	OpTypeFusedGelu
	OpTypeFusedDense
	OpTypeFusedScaledDotProductAttention
	OpTypeFusedScaledDotProductAttentionVJP
	OpTypeFusedAttentionQKVProjection
	OpTypeFusedQuantizedDense
	OpTypeQuantizedEmbeddingLookup

	// OpTypeLast should always be kept the last, it is used as a counter/marker for OpType.
	OpTypeLast
)