Skip to content

Commit

Permalink
fix!: OpDef serialisation (#1013)
Browse files Browse the repository at this point in the history 
BREAKING CHANGE:
* OpRef serialisation schema

---------

Co-authored-by: Mark Koch <48097969+mark-koch@users.noreply.github.com>
  • Loading branch information
@doug-q @mark-koch
doug-q and mark-koch committed May 10, 2024
1 parent 33ab50d commit 3d8f6f6
Show file tree
Hide file tree
Showing 7 changed files with 210 additions and 22 deletions.
18 changes: 10 additions & 8 deletions hugr-py/src/hugr/serialization/ops.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -500,18 +500,20 @@ class Config:
# --------------------------------------


class OpDef(BaseOp, populate_by_name=True):
class FixedHugr(ConfiguredBaseModel):
extensions: ExtensionSet
hugr: Any


class OpDef(ConfiguredBaseModel, populate_by_name=True):
"""Serializable definition for dynamically loaded operations."""

extension: ExtensionId
name: str # Unique identifier of the operation.
description: str # Human readable description of the operation.
inputs: list[tuple[str | None, Type]]
outputs: list[tuple[str | None, Type]]
misc: dict[str, Any] # Miscellaneous data associated with the operation.
def_: str | None = Field(
..., alias="def"
) # (YAML?)-encoded definition of the operation.
extension_reqs: ExtensionSet # Resources required to execute this operation.
misc: dict[str, Any] | None = None
signature: PolyFuncType | None = None
lower_funcs: list[FixedHugr]


# Now that all classes are defined, we need to update the ForwardRefs in all type
Expand Down
3 changes: 2 additions & 1 deletion hugr-py/src/hugr/serialization/testing_hugr.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from pydantic import ConfigDict
from typing import Literal
from .tys import Type, SumType, PolyFuncType, ConfiguredBaseModel, model_rebuild
from .ops import Value, OpType, classes as ops_classes
from .ops import Value, OpType, OpDef, classes as ops_classes


class TestingHugr(ConfiguredBaseModel):
Expand All @@ -14,6 +14,7 @@ class TestingHugr(ConfiguredBaseModel):
poly_func_type: PolyFuncType | None = None
value: Value | None = None
optype: OpType | None = None
op_def: OpDef | None = None

@classmethod
def get_version(cls) -> str:
Expand Down
2 changes: 1 addition & 1 deletion hugr/src/extension.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ pub use infer::{ExtensionSolution, InferExtensionError};

mod op_def;
pub use op_def::{
CustomSignatureFunc, CustomValidator, OpDef, SignatureFromArgs, SignatureFunc,
CustomSignatureFunc, CustomValidator, LowerFunc, OpDef, SignatureFromArgs, SignatureFunc,
ValidateJustArgs, ValidateTypeArgs,
};
mod type_def;
Expand Down
31 changes: 20 additions & 11 deletions hugr/src/extension/op_def.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ pub struct CustomValidator {
#[serde(flatten)]
poly_func: PolyFuncType,
#[serde(skip)]
validate: Box<dyn ValidateTypeArgs>,
pub(crate) validate: Box<dyn ValidateTypeArgs>,
}

impl CustomValidator {
Expand Down Expand Up @@ -265,11 +265,19 @@ impl Debug for SignatureFunc {
/// Different ways that an [OpDef] can lower operation nodes i.e. provide a Hugr
/// that implements the operation using a set of other extensions.
#[derive(serde::Deserialize, serde::Serialize)]
#[serde(untagged)]
pub enum LowerFunc {
/// Lowering to a fixed Hugr. Since this cannot depend upon the [TypeArg]s,
/// this will generally only be applicable if the [OpDef] has no [TypeParam]s.
#[serde(rename = "hugr")]
FixedHugr(ExtensionSet, Hugr),
FixedHugr {
/// The extensions required by the [`Hugr`]
extensions: ExtensionSet,
/// The [`Hugr`] to be used to replace [CustomOp]s matching the parent
/// [OpDef]
///
/// [CustomOp]: crate::ops::CustomOp
hugr: Hugr,
},
/// Custom binary function that can (fallibly) compute a Hugr
/// for the particular instance and set of available extensions.
#[serde(skip)]
Expand All @@ -279,7 +287,7 @@ pub enum LowerFunc {
impl Debug for LowerFunc {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::FixedHugr(_, _) => write!(f, "FixedHugr"),
Self::FixedHugr { .. } => write!(f, "FixedHugr"),
Self::CustomFunc(_) => write!(f, "<custom lower>"),
}
}
Expand All @@ -305,8 +313,7 @@ pub struct OpDef {
signature_func: SignatureFunc,
// Some operations cannot lower themselves and tools that do not understand them
// can only treat them as opaque/black-box ops.
#[serde(flatten)]
lower_funcs: Vec<LowerFunc>,
pub(crate) lower_funcs: Vec<LowerFunc>,

/// Operations can optionally implement [`ConstFold`] to implement constant folding.
#[serde(skip)]
Expand Down Expand Up @@ -360,9 +367,9 @@ impl OpDef {
self.lower_funcs
.iter()
.flat_map(|f| match f {
LowerFunc::FixedHugr(req_res, h) => {
if available_extensions.is_superset(req_res) {
Some(h.clone())
LowerFunc::FixedHugr { extensions, hugr } => {
if available_extensions.is_superset(extensions) {
Some(hugr.clone())
} else {
None
}
Expand Down Expand Up @@ -477,7 +484,6 @@ mod test {
use crate::std_extensions::collections::{EXTENSION, LIST_TYPENAME};
use crate::types::Type;
use crate::types::{type_param::TypeParam, FunctionType, PolyFuncType, TypeArg, TypeBound};
use crate::Hugr;
use crate::{const_extension_ids, Extension};

const_extension_ids! {
Expand All @@ -495,7 +501,10 @@ mod test {
let type_scheme = PolyFuncType::new(vec![TP], FunctionType::new_endo(vec![list_of_var]));

let def = e.add_op(OP_NAME, "desc".into(), type_scheme)?;
def.add_lower_func(LowerFunc::FixedHugr(ExtensionSet::new(), Hugr::default()));
def.add_lower_func(LowerFunc::FixedHugr {
extensions: ExtensionSet::new(),
hugr: crate::builder::test::simple_dfg_hugr(), // this is nonsense, but we are not testing the actual lowering here
});
def.add_misc("key", Default::default());
assert_eq!(def.description(), "desc");
assert_eq!(def.lower_funcs.len(), 1);
Expand Down
2 changes: 1 addition & 1 deletion hugr/src/hugr/serialize/test.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -377,7 +377,7 @@ fn roundtrip_sumtype(#[case] sum_type: SumType) {
// #[case(Value::extension(ConstF64::new(std::f64::NAN)))]
// #[case(Value::extension(ConstF64::new(std::f64::INFINITY)))]
// #[case(Value::extension(ConstF64::new(std::f64::NEG_INFINITY)))]
#[case(Value::extension(ConstF64::new(std::f64::MIN_POSITIVE)))]
#[case(Value::extension(ConstF64::new(f64::MIN_POSITIVE)))]
#[case(Value::sum(1,[Value::extension(ConstInt::new_u(2,1).unwrap())], SumType::new([vec![], vec![INT_TYPES[2].clone()]])).unwrap())]
#[case(Value::tuple([Value::false_val(), Value::extension(ConstInt::new_s(2,1).unwrap())]))]
#[case(Value::function(crate::builder::test::simple_dfg_hugr()).unwrap())]
Expand Down
88 changes: 88 additions & 0 deletions specification/schema/testing_hugr_schema_strict_v1.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -817,6 +817,27 @@
"title": "ExtensionsParam",
"type": "object"
},
"FixedHugr": {
"additionalProperties": false,
"properties": {
"extensions": {
"items": {
"type": "string"
},
"title": "Extensions",
"type": "array"
},
"hugr": {
"title": "Hugr"
}
},
"required": [
"extensions",
"hugr"
],
"title": "FixedHugr",
"type": "object"
},
"FuncDecl": {
"additionalProperties": false,
"description": "External function declaration, linked at runtime.",
Expand Down Expand Up @@ -1363,6 +1384,62 @@
"title": "Noop",
"type": "object"
},
"OpDef": {
"additionalProperties": false,
"description": "Serializable definition for dynamically loaded operations.",
"properties": {
"extension": {
"title": "Extension",
"type": "string"
},
"name": {
"title": "Name",
"type": "string"
},
"description": {
"title": "Description",
"type": "string"
},
"misc": {
"anyOf": [
{
"type": "object"
},
{
"type": "null"
}
],
"default": null,
"title": "Misc"
},
"signature": {
"anyOf": [
{
"$ref": "#/$defs/PolyFuncType"
},
{
"type": "null"
}
],
"default": null
},
"lower_funcs": {
"items": {
"$ref": "#/$defs/FixedHugr"
},
"title": "Lower Funcs",
"type": "array"
}
},
"required": [
"extension",
"name",
"description",
"lower_funcs"
],
"title": "OpDef",
"type": "object"
},
"OpType": {
"description": "A constant operation.",
"discriminator": {
Expand Down Expand Up @@ -2325,6 +2402,17 @@
}
],
"default": null
},
"op_def": {
"anyOf": [
{
"$ref": "#/$defs/OpDef"
},
{
"type": "null"
}
],
"default": null
}
},
"title": "TestingHugr",
Expand Down
88 changes: 88 additions & 0 deletions specification/schema/testing_hugr_schema_v1.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -817,6 +817,27 @@
"title": "ExtensionsParam",
"type": "object"
},
"FixedHugr": {
"additionalProperties": true,
"properties": {
"extensions": {
"items": {
"type": "string"
},
"title": "Extensions",
"type": "array"
},
"hugr": {
"title": "Hugr"
}
},
"required": [
"extensions",
"hugr"
],
"title": "FixedHugr",
"type": "object"
},
"FuncDecl": {
"additionalProperties": true,
"description": "External function declaration, linked at runtime.",
Expand Down Expand Up @@ -1363,6 +1384,62 @@
"title": "Noop",
"type": "object"
},
"OpDef": {
"additionalProperties": true,
"description": "Serializable definition for dynamically loaded operations.",
"properties": {
"extension": {
"title": "Extension",
"type": "string"
},
"name": {
"title": "Name",
"type": "string"
},
"description": {
"title": "Description",
"type": "string"
},
"misc": {
"anyOf": [
{
"type": "object"
},
{
"type": "null"
}
],
"default": null,
"title": "Misc"
},
"signature": {
"anyOf": [
{
"$ref": "#/$defs/PolyFuncType"
},
{
"type": "null"
}
],
"default": null
},
"lower_funcs": {
"items": {
"$ref": "#/$defs/FixedHugr"
},
"title": "Lower Funcs",
"type": "array"
}
},
"required": [
"extension",
"name",
"description",
"lower_funcs"
],
"title": "OpDef",
"type": "object"
},
"OpType": {
"description": "A constant operation.",
"discriminator": {
Expand Down Expand Up @@ -2325,6 +2402,17 @@
}
],
"default": null
},
"op_def": {
"anyOf": [
{
"$ref": "#/$defs/OpDef"
},
{
"type": "null"
}
],
"default": null
}
},
"title": "TestingHugr",
Expand Down

0 comments on commit 3d8f6f6

Please sign in to comment.