Sketching operations

SHerLOC/AST/Basic.lean

@@ -0,0 +1,11 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Jean-Baptiste Tristan
+-/
+import SHerLOC.AST.Types
+import SHerLOC.AST.Identifiers
+import SHerLOC.AST.Constants
+import SHerLOC.AST.Operations
+import SHerLOC.AST.Functions
+import SHerLOC.AST.Programs

SHerLOC/AST/Constants.lean

@@ -0,0 +1,25 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Jean-Baptiste Tristan
+-/
+import SHerLOC.AST.Types
+
+/-!
+# Constants
+
+-/
+
+namespace STableHLO
+
+inductive Constant where
+  | booleanConstant (literal : Bool)
+  | integerConstant (literal : Int) (type : IntegerType)
+  | floatConstant (literal : Float) (type : FloatType)
+  | complexConstant (real imaginary : ComplexType)
+  | tensorConstant
+  | quantizedTensorConstant
+  | stringConstant (literal : String)
+  | enumConstant (choices : List String)
+
+end STableHLO

SHerLOC/AST/Functions.lean

@@ -0,0 +1,23 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Jean-Baptiste Tristan
+-/
+import SHerLOC.AST.Types
+import SHerLOC.AST.Identifiers
+import SHerLOC.AST.Operations
+
+/-!
+# Functions
+
+-/
+
+namespace StableHLO
+
+structure Function where
+  funcId : FuncId
+  funcInputs : List (ValueId × ValueType)
+  funcOutputs : List ValueType
+  funcBody : List Operation
+
+end StableHLO

SHerLOC/Identifiers.lean → SHerLOC/AST/Identifiers.lean

@@ -9,5 +9,10 @@ Authors: Jean-Baptiste Tristan
 
 -/
 
-def FuncID := Nat
-def ValueID := Nat
+namespace StableHLO
+
+def FuncId := Nat
+def ValueId := Nat
+def UnusedId := Nat
+
+end StableHLO

SHerLOC/AST/Operations.lean

@@ -0,0 +1,149 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Jean-Baptiste Tristan
+-/
+import SHerLOC.AST.Constants
+import SHerLOC.AST.Identifiers
+import SHerLOC.AST.Types
+
+/-!
+# Operations
+
+-/
+
+namespace StableHLO
+
+inductive OpName where
+  | abs
+  | add
+  | after_all
+  | all_gather
+  | all_reduce
+  | all_to_all
+  | and
+  | atan2
+  | batch_norm_grad
+  | batch_norm_inference
+  | batch_norm_training
+  | bitcast_convert
+  | broadcast_in_dim
+  | case
+  | cbrt
+  | ceil
+  | cholesky
+  | clamp
+  | collective_broadcast
+  | collective_permute
+  | compare
+  | complex
+  | composite
+  | concatenate
+  | constant
+  | convert
+  | convolution
+  | cosine
+  | count_leading_zeros
+  | custom_call
+  | divide
+  | dot_general
+  | dynamic_broadcast_in_dim
+  | dynamic_conv
+  | dynamic_gather
+  | dynamic_iota
+  | dynamic_pad
+  | dynamic_reshape
+  | dynamic_slice
+  | dynamic_update_slice
+  | exponential
+  | exponential_minus_one
+  | fft
+  | floor
+  | gather
+  | get_dimension_size
+  | get_tuple_element
+  | if
+  | imag
+  | infeed
+  | iota
+  | is_finite
+  | log
+  | log_plus_one
+  | logistic
+  | map
+  | maximum
+  | minimum
+  | multiply
+  | negate
+  | not
+  | optimization_barrier
+  | or
+  | outfeed
+  | pad
+  | partition_id
+  | popcnt
+  | power
+  | real
+  | recv
+  | reduce
+  | reduce_precision
+  | reduce_scatter
+  | reduce_window
+  | remainder
+  | replica_id
+  | reshape
+  | reverse
+  | rng
+  | rng_bit_generator
+  | round_nearest_afz
+  | round_nearest_even
+  | rsqrt
+  | scatter
+  | select
+  | select_and_scatter
+  | send
+  | shift_left
+  | shift_right_arithmetic
+  | shift_right_logical
+  | sign
+  | sine
+  | slice
+  | sort
+  | sqrt
+  | subtract
+  | tan
+  | tanh
+  | transpose
+  | triangular_solve
+  | tuple
+  | uniform_dequantize
+  | uniform_quantize
+  | while
+  | xor
+
+mutual
+
+  inductive InputFunc where
+    | mk
+      (id : UnusedId)
+      (funcInputs : List ValueId)
+      (body : List Operation)
+
+  inductive Operation where
+    | stable
+      (name : OpName)(inputValues : List ValueId)
+      (inputFunctions : List InputFunc)
+      (inputAttributes : List Constant)
+      (outputs : List ValueId)
+      (signature : FunctionType)
+    | return
+      (operands : ValueId)
+      (signature : ValueType)
+    | call
+      (callee : FuncId)
+      (arguments : ValueId)
+      (signature : FunctionType)
+
+end
+
+end StableHLO

SHerLOC/AST/Programs.lean

@@ -0,0 +1,17 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Jean-Baptiste Tristan
+-/
+import SHerLOC.AST.Functions
+
+/-!
+# Programs
+
+-/
+
+namespace StableHLO
+
+def Program := List Function
+
+end StableHLO

SHerLOC/AST/Types.lean

@@ -0,0 +1,77 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Jean-Baptiste Tristan
+-/
+
+/-!
+# Types
+
+-/
+
+namespace StableHLO
+
+inductive Signedness where
+  | signed
+  | unsigned
+
+inductive IntegerSize where
+  | b2
+  | b4
+  | b8
+  | b16
+  | b32
+  | b64
+
+structure IntegerType where
+  sign : Signedness
+  size : IntegerSize
+
+inductive FloatType where
+  | f8E4M3FN
+  | f8E5M2
+  | f8E4M3FNUZ
+  | f8E5M2FNUZ
+  | f8E4M3B11FNUZ
+  | bf16
+  | f16
+  | f32
+  | f64
+  | tf32
+
+inductive ComplexType where
+  | f32
+  | f64
+
+inductive TensorElementType where
+  | booleanType
+  | integerType (t : IntegerType)
+  | floatType (t: FloatType)
+  | complexType (t: ComplexType)
+
+inductive QuantizedTensorElementType where
+  | quant : Signedness → IntegerSize → Int → Int → FloatSize → Int → List (Float × Int) → QuantizedTensorElementType
+
+inductive ValueType where
+  | tensorType (shape : List Int) (typ : TensorElementType)
+  | quantizedTensorType (shape : List Int) (typ : QuantizedTensorElementType)
+  | tokenType
+  | tupleType (elements : List Valuetype)
+
+inductive StringType where
+
+structure FunctionType where
+  domain : List ValueType
+  range : List ValueType
+
+inductive NonValueType where
+  | tensorElementType (t : TensorElementType)
+  | quantizedTensorElementType (t: QuantizedTensorElementType)
+  | functionType (t : FunctionType)
+  | stringType (t : StringType)
+
+inductive SType where
+  | valueType (t : ValueType)
+  | nonValueType (t : NonValueType)
+
+end StableHLO

SHerLOC/Basic.lean

@@ -1,3 +1,6 @@
-import SHerLOC.Types
-import SHerLOC.Identifiers
-import SHerLOC.Constants
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Jean-Baptiste Tristan
+-/
+import SHerLOC.AST.Basic