From c60402ca35816bc7aac0fd473f1acd2b29f8eb2a Mon Sep 17 00:00:00 2001 From: John Tristan Date: Tue, 20 Aug 2024 18:11:32 -0400 Subject: [PATCH] Sketching operations (#5) --- SHerLOC/AST/Basic.lean | 11 +++ SHerLOC/AST/Constants.lean | 25 +++++ SHerLOC/AST/Functions.lean | 23 +++++ SHerLOC/{ => AST}/Identifiers.lean | 9 +- SHerLOC/AST/Operations.lean | 149 +++++++++++++++++++++++++++++ SHerLOC/AST/Programs.lean | 17 ++++ SHerLOC/AST/Types.lean | 77 +++++++++++++++ SHerLOC/Basic.lean | 9 +- SHerLOC/Constants.lean | 21 ---- SHerLOC/Types.lean | 62 ------------ 10 files changed, 315 insertions(+), 88 deletions(-) create mode 100644 SHerLOC/AST/Basic.lean create mode 100644 SHerLOC/AST/Constants.lean create mode 100644 SHerLOC/AST/Functions.lean rename SHerLOC/{ => AST}/Identifiers.lean (68%) create mode 100644 SHerLOC/AST/Operations.lean create mode 100644 SHerLOC/AST/Programs.lean create mode 100644 SHerLOC/AST/Types.lean delete mode 100644 SHerLOC/Constants.lean delete mode 100644 SHerLOC/Types.lean diff --git a/SHerLOC/AST/Basic.lean b/SHerLOC/AST/Basic.lean new file mode 100644 index 0000000..9ee343e --- /dev/null +++ b/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 diff --git a/SHerLOC/AST/Constants.lean b/SHerLOC/AST/Constants.lean new file mode 100644 index 0000000..a5df97b --- /dev/null +++ b/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 diff --git a/SHerLOC/AST/Functions.lean b/SHerLOC/AST/Functions.lean new file mode 100644 index 0000000..9861b71 --- /dev/null +++ b/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 diff --git a/SHerLOC/Identifiers.lean b/SHerLOC/AST/Identifiers.lean similarity index 68% rename from SHerLOC/Identifiers.lean rename to SHerLOC/AST/Identifiers.lean index e87d23b..fbee4eb 100644 --- a/SHerLOC/Identifiers.lean +++ b/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 diff --git a/SHerLOC/AST/Operations.lean b/SHerLOC/AST/Operations.lean new file mode 100644 index 0000000..782ebc4 --- /dev/null +++ b/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 diff --git a/SHerLOC/AST/Programs.lean b/SHerLOC/AST/Programs.lean new file mode 100644 index 0000000..5ad83cf --- /dev/null +++ b/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 diff --git a/SHerLOC/AST/Types.lean b/SHerLOC/AST/Types.lean new file mode 100644 index 0000000..8ff0b15 --- /dev/null +++ b/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 diff --git a/SHerLOC/Basic.lean b/SHerLOC/Basic.lean index 8430dc8..ca1ed6a 100644 --- a/SHerLOC/Basic.lean +++ b/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 diff --git a/SHerLOC/Constants.lean b/SHerLOC/Constants.lean deleted file mode 100644 index 6e6a94e..0000000 --- a/SHerLOC/Constants.lean +++ /dev/null @@ -1,21 +0,0 @@ -/- -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.Types - -/-! -# Constants - --/ - -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 diff --git a/SHerLOC/Types.lean b/SHerLOC/Types.lean deleted file mode 100644 index 7d60fc1..0000000 --- a/SHerLOC/Types.lean +++ /dev/null @@ -1,62 +0,0 @@ -/- -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 - --/ - -inductive Signedness where - | Signed - | Unsigned - -inductive IntegerSize where - | b2 - | b4 - | b8 - | b16 - | b32 - | b64 - -inductive IntegerType where - | IntegerType (sign : Signedness) (size : IntegerSize) - -inductive FloatType where - | f8E4M3FN - | f8E5M2 - | f8E4M3FNUZ - | f8E5M2FNUZ - | f8E4M3B11FNUZ - | bf16 - | f16 - | f32 - | f64 - -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 : List Int → TensorElementType → ValueType - | QuantizedTensorType : List Int → QuantizedTensorElementType → ValueType - | TokenType - | TupleType : List Valuetype → ValueType - -inductive TensorFloar32 where - -inductive StringType where - -inductive FunctionType where - | FunctionType : List ValueType → List ValueType → FunctionType