Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

dialects: (stencil) add StencilType constraints #3254

Merged
merged 3 commits into from
Oct 8, 2024
Merged

dialects: (stencil) add StencilType constraints #3254

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ea0dee5
dialects: (stencil) add StencilTypeConstr
superlopuh Sep 17, 2024
74f2106
revert comment change
superlopuh Oct 7, 2024
45f4e92
add constr helper
superlopuh Oct 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions xdsl/dialects/csl/csl_stencil.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,7 @@ class PrefetchOp(IRDLOperation):
name = "csl_stencil.prefetch"

input_stencil = operand_def(
base(stencil.StencilType[Attribute]) | AnyMemRefTypeConstr | AnyTensorTypeConstr
stencil.StencilTypeConstr | AnyMemRefTypeConstr | AnyTensorTypeConstr
)

swaps = prop_def(builtin.ArrayAttr[ExchangeDeclarationAttr])
Expand Down Expand Up @@ -202,12 +202,12 @@ class ApplyOp(IRDLOperation):

name = "csl_stencil.apply"

field = operand_def(base(stencil.StencilType[Attribute]) | AnyMemRefTypeConstr)
field = operand_def(stencil.StencilTypeConstr | AnyMemRefTypeConstr)

accumulator = operand_def(AnyTensorTypeConstr | AnyMemRefTypeConstr)

args = var_operand_def(Attribute)
dest = var_operand_def(base(stencil.FieldType[Attribute]) | AnyMemRefTypeConstr)
dest = var_operand_def(stencil.FieldTypeConstr | AnyMemRefTypeConstr)

receive_chunk = region_def()
done_exchange = region_def()
Expand All @@ -220,7 +220,7 @@ class ApplyOp(IRDLOperation):

bounds = opt_prop_def(stencil.StencilBoundsAttr)

res = var_result_def(stencil.StencilType)
res = var_result_def(stencil.StencilTypeConstr)

traits = frozenset(
[
Expand Down Expand Up @@ -376,7 +376,7 @@ def get_rank(self) -> int:
res_type = self.dest[0].type
else:
res_type = self.res[0].type
if isattr(res_type, base(stencil.StencilType[Attribute])):
if isattr(res_type, stencil.StencilTypeConstr):
return res_type.get_num_dims()
elif self.bounds:
return len(self.bounds.ub)
Expand Down Expand Up @@ -423,7 +423,7 @@ class AccessOp(IRDLOperation):

name = "csl_stencil.access"
op = operand_def(
AnyMemRefTypeConstr | base(stencil.StencilType[Attribute]) | AnyTensorTypeConstr
AnyMemRefTypeConstr | stencil.StencilTypeConstr | AnyTensorTypeConstr
)
offset = prop_def(stencil.IndexAttr)
offset_mapping = opt_prop_def(stencil.IndexAttr)
Expand Down Expand Up @@ -502,7 +502,7 @@ def parse(cls, parser: Parser):
props["offset_mapping"] = stencil.IndexAttr.get(*offset_mapping)
parser.parse_punctuation(":")
res_type = parser.parse_attribute()
if isattr(res_type, base(stencil.StencilType[Attribute])):
if isattr(res_type, stencil.StencilTypeConstr):
return cls.build(
operands=[temp],
result_types=[res_type.get_element_type()],
Expand Down Expand Up @@ -535,7 +535,7 @@ def verify_(self) -> None:
raise VerifyException(
f"{type(self)} access to own data requires{self.op.type} but found {self.result.type}"
)
elif isattr(self.op.type, base(stencil.StencilType[Attribute])):
elif isattr(self.op.type, stencil.StencilTypeConstr):
if not self.result.type == self.op.type.get_element_type():
raise VerifyException(
f"{type(self)} access to own data requires{self.op.type.get_element_type()} but found {self.result.type}"
Expand Down
2 changes: 1 addition & 1 deletion xdsl/dialects/experimental/dmp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -686,7 +686,7 @@ class SwapOp(IRDLOperation):

name = "dmp.swap"

input_stencil = operand_def(stencil.StencilType[Attribute])
input_stencil = operand_def(stencil.StencilTypeConstr)
swapped_values = opt_result_def(stencil.TempType[Attribute])

swaps = attr_def(builtin.ArrayAttr[ExchangeDeclarationAttr])
Expand Down
196 changes: 86 additions & 110 deletions xdsl/dialects/stencil.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,11 @@
)
from xdsl.irdl import (
AnyAttr,
AnyOf,
AttrSizedOperandSegments,
BaseAttr,
ConstraintContext,
ConstraintVar,
GenericAttrConstraint,
IRDLOperation,
MessageConstraint,
ParamAttrConstraint,
Expand Down Expand Up @@ -70,6 +71,7 @@
)
from xdsl.utils.exceptions import VerifyException
from xdsl.utils.hints import isa
from xdsl.utils.isattr import isattr

_FieldTypeElement = TypeVar("_FieldTypeElement", bound=Attribute, covariant=True)

Expand Down Expand Up @@ -357,6 +359,20 @@ def __init__(
nbounds = bounds
return super().__init__([nbounds, element_type])

@classmethod
def constr(
cls,
*,
bounds: GenericAttrConstraint[Attribute] | None = None,
element_type: GenericAttrConstraint[_FieldTypeElement] | None = None,
) -> (
BaseAttr[StencilType[Attribute]]
| ParamAttrConstraint[StencilType[_FieldTypeElement]]
):
if bounds is None and element_type is None:
return BaseAttr(cls)
return ParamAttrConstraint(cls, (bounds, element_type))


@irdl_attr_definition
class FieldType(
Expand Down Expand Up @@ -390,6 +406,10 @@ class TempType(
name = "stencil.temp"


StencilTypeConstr = StencilType[Attribute].constr()
FieldTypeConstr = FieldType[Attribute].constr()
TempTypeConstr = TempType[Attribute].constr()

AnyTempType: TypeAlias = TempType[Attribute]


Expand Down Expand Up @@ -696,27 +716,19 @@ class CastOp(IRDLOperation):
name = "stencil.cast"

field = operand_def(
ParamAttrConstraint(
FieldType,
[
Attribute,
MessageConstraint(
VarConstraint("T", AnyAttr()),
"Input and output fields must have the same element types",
),
],
FieldType[Attribute].constr(
element_type=MessageConstraint(
VarConstraint("T", AnyAttr()),
"Input and output fields must have the same element types",
)
)
)
result = result_def(
ParamAttrConstraint(
FieldType,
[
Attribute,
MessageConstraint(
VarConstraint("T", AnyAttr()),
"Input and output fields must have the same element types",
),
],
FieldType[Attribute].constr(
element_type=MessageConstraint(
VarConstraint("T", AnyAttr()),
"Input and output fields must have the same element types",
)
)
)

Expand Down Expand Up @@ -851,17 +863,14 @@ class DynAccessOp(IRDLOperation):
T = Annotated[Attribute, ConstraintVar("T")]

temp = operand_def(
ParamAttrConstraint(
StencilType,
[
Attribute,
MessageConstraint(
VarConstraint("T", AnyAttr()),
"Expected result type to be the accessed temp's element type.",
),
],
StencilType[Attribute].constr(
element_type=MessageConstraint(
VarConstraint("T", AnyAttr()),
"Expected result type to be the accessed temp's element type.",
)
)
)

offset = var_operand_def(builtin.IndexType())
lb = attr_def(IndexAttr)
ub = attr_def(IndexAttr)
Expand Down Expand Up @@ -1007,15 +1016,11 @@ class AccessOp(IRDLOperation):

name = "stencil.access"
temp = operand_def(
ParamAttrConstraint(
StencilType,
[
Attribute,
MessageConstraint(
VarConstraint("T", AnyAttr()),
"Expected return type to match the accessed temp's element type.",
),
],
StencilType[Attribute].constr(
element_type=MessageConstraint(
VarConstraint("T", AnyAttr()),
"Expected return type to match the accessed temp's element type.",
)
)
)
offset = attr_def(IndexAttr)
Expand Down Expand Up @@ -1094,7 +1099,7 @@ def parse(cls, parser: Parser):
attrs["offset_mapping"] = IndexAttr.get(*offset_mapping)
parser.parse_punctuation(":")
res_type = parser.parse_attribute()
if not isa(res_type, StencilType[Attribute]):
if not isattr(res_type, StencilTypeConstr):
parser.raise_error(
"Expected return type to be a stencil.temp or stencil.field"
)
Expand Down Expand Up @@ -1142,7 +1147,7 @@ def verify_(self) -> None:
apply.verify_()

temp_type = self.temp.type
assert isa(temp_type, StencilType[Attribute])
assert isattr(temp_type, StencilTypeConstr)
if temp_type.get_num_dims() != apply.get_rank():
if self.offset_mapping is None:
raise VerifyException(
Expand Down Expand Up @@ -1226,25 +1231,18 @@ class LoadOp(IRDLOperation):
T = Annotated[Attribute, ConstraintVar("T")]

field = operand_def(
ParamAttrConstraint(
FieldType,
[
StencilBoundsAttr,
MessageConstraint(
VarConstraint("T", AnyAttr()), "Expected element types to match."
),
],
FieldType[Attribute].constr(
bounds=base(StencilBoundsAttr),
element_type=MessageConstraint(
VarConstraint("T", AnyAttr()), "Expected element types to match."
),
)
)
res = result_def(
ParamAttrConstraint(
TempType,
[
Attribute,
MessageConstraint(
VarConstraint("T", AnyAttr()), "Expected element types to match."
),
],
TempType[Attribute].constr(
element_type=MessageConstraint(
VarConstraint("T", AnyAttr()), "Expected element types to match."
)
)
)

Expand Down Expand Up @@ -1309,33 +1307,27 @@ class BufferOp(IRDLOperation):
T = Annotated[TempType[_FieldTypeElement], ConstraintVar("T")]

temp = operand_def(
ParamAttrConstraint(
TempType,
[
MessageConstraint(
VarConstraint("B", AnyAttr()),
"Expected input and output to have the same bounds",
),
MessageConstraint(
VarConstraint("E", AnyAttr()),
"Expected input and output to have the same element type",
),
],
TempType[Attribute].constr(
bounds=MessageConstraint(
VarConstraint("B", AnyAttr()),
"Expected input and output to have the same bounds",
),
element_type=MessageConstraint(
VarConstraint("E", AnyAttr()),
"Expected input and output to have the same element type",
),
)
)
res = result_def(
ParamAttrConstraint(
StencilType,
[
MessageConstraint(
VarConstraint("B", AnyAttr()),
"Expected input and output to have the same bounds",
),
MessageConstraint(
VarConstraint("E", AnyAttr()),
"Expected input and output to have the same element type",
),
],
StencilType[Attribute].constr(
bounds=MessageConstraint(
VarConstraint("B", AnyAttr()),
"Expected input and output to have the same bounds",
),
element_type=MessageConstraint(
VarConstraint("E", AnyAttr()),
"Expected input and output to have the same element type",
),
)
)

Expand Down Expand Up @@ -1414,40 +1406,24 @@ class StoreOp(IRDLOperation):
name = "stencil.store"

temp = operand_def(
ParamAttrConstraint(
TempType,
[
Attribute,
MessageConstraint(
AnyOf(
[
VarConstraint("T", AnyAttr()),
TensorIgnoreSizeConstraint("T", AnyAttr()),
]
),
"Input and output fields must have the same element types",
),
],
TempType[Attribute].constr(
element_type=MessageConstraint(
VarConstraint("T", AnyAttr())
| TensorIgnoreSizeConstraint("T", AnyAttr()),
"Input and output fields must have the same element types",
),
)
)
# field = operand_def(FieldType)
field = operand_def(
ParamAttrConstraint(
FieldType,
[
MessageConstraint(
StencilBoundsAttr, "Output type's size must be explicit"
),
MessageConstraint(
AnyOf(
[
VarConstraint("T", AnyAttr()),
TensorIgnoreSizeConstraint("T", AnyAttr()),
]
),
"Input and output fields must have the same element types",
),
],
FieldType[Attribute].constr(
bounds=MessageConstraint(
StencilBoundsAttr, "Output type's size must be explicit"
),
element_type=MessageConstraint(
VarConstraint("T", AnyAttr())
| TensorIgnoreSizeConstraint("T", AnyAttr()),
"Input and output fields must have the same element types",
),
)
)
bounds = attr_def(StencilBoundsAttr)
Expand Down
7 changes: 5 additions & 2 deletions xdsl/irdl/constraints.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -361,10 +361,13 @@ def __init__(
self,
base_attr: type[ParametrizedAttributeCovT],
param_constrs: Sequence[
(Attribute | type[Attribute] | GenericAttrConstraint[Attribute])
(Attribute | type[Attribute] | GenericAttrConstraint[Attribute] | None)
],
):
constrs = tuple(attr_constr_coercion(constr) for constr in param_constrs)
constrs = tuple(
attr_constr_coercion(constr) if constr is not None else AnyAttr()
for constr in param_constrs
)
object.__setattr__(self, "base_attr", base_attr)
object.__setattr__(self, "param_constrs", constrs)

Expand Down
2 changes: 1 addition & 1 deletion xdsl/transforms/convert_stencil_to_csl_stencil.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,7 +206,7 @@ def match_and_rewrite(self, op: dmp.SwapOp, rewriter: PatternRewriter, /):

assert isattr(
op.input_stencil.type,
AnyMemRefTypeConstr | base(stencil.StencilType[Attribute]),
AnyMemRefTypeConstr | stencil.StencilTypeConstr,
)
assert isa(
t_type := op.input_stencil.type.get_element_type(), TensorType[Attribute]
Expand Down
Loading
Loading