# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Models DAGs of scalar math expressions.
Used for generating region bodies at the "math" level where they are still type
polymorphic. This is modeled to be polymorphic by attribute name for interop
with serialization schemes that are just plain-old-dicts.
These classes are typically not user accessed and are created as a by-product
of interpreting a comprehension DSL and model the operations to perform in the
op body. The class hierarchy is laid out to map well to a form of YAML that
can be easily consumed from the C++ side, not necessarily for ergonomics.
"""
from typing import Optional, Sequence
from .comprehension import *
from .types import *
from .yaml_helper import *
__all__ = [
"ScalarAssign",
"ScalarFn",
"ScalarArg",
"ScalarConst",
"ScalarIndex",
"ScalarExpression",
]
class ScalarFn:
"""A type of ScalarExpression that applies a function."""
def __init__(
self,
kind: "FunctionKind",
fn_name: Optional[str],
attr_name: Optional[str],
type_var: Optional["TypeVar"],
operands: Sequence["ScalarExpression"],
):
if bool(fn_name) + bool(attr_name) != 1:
raise ValueError("One of 'fn_name', 'attr_name' must be specified")
self.kind = kind
self.fn_name = fn_name
self.attr_name = attr_name
self.type_var = type_var
self.operands = operands
def expr(self) -> "ScalarExpression":
return ScalarExpression(scalar_fn=self)
def __repr__(self):
name = self.fn_name if self.fn_name else self.attr_name
return (
f"ScalarFn<{self.kind.name}.{name}>(type_var={self.type_var}, "
f"operands=[{', '.join(self.operands)}])"
)
class ScalarArg:
"""A type of ScalarExpression that references a named argument."""
def __init__(self, arg: str):
self.arg = arg
def expr(self) -> "ScalarExpression":
return ScalarExpression(scalar_arg=self)
def __repr__(self):
return f"(ScalarArg({self.arg})"
class ScalarConst:
"""A type of ScalarExpression representing a constant."""
def __init__(self, value: str):
self.value = value
def expr(self) -> "ScalarExpression":
return ScalarExpression(scalar_const=self)
def __repr__(self):
return f"(ScalarConst({self.value})"
class ScalarIndex:
"""A type of ScalarExpression accessing an iteration index."""
def __init__(self, dim: int):
self.dim = dim
def expr(self) -> "ScalarExpression":
return ScalarExpression(scalar_index=self)
def __repr__(self):
return f"(ScalarIndex({self.dim})"
class ScalarExpression(YAMLObject):
"""An expression on scalar values.
Can be one of:
- ScalarFn
- ScalarArg
- ScalarConst
- ScalarIndex
"""
yaml_tag = "!ScalarExpression"
def __init__(
self,
scalar_fn: Optional[ScalarFn] = None,
scalar_arg: Optional[ScalarArg] = None,
scalar_const: Optional[ScalarConst] = None,
scalar_index: Optional[ScalarIndex] = None,
):
if (
bool(scalar_fn) + bool(scalar_arg) + bool(scalar_const) + bool(scalar_index)
) != 1:
raise ValueError(
"One of 'scalar_fn', 'scalar_arg', 'scalar_const', or "
"'scalar_index' must be specified"
)
self.scalar_fn = scalar_fn
self.scalar_arg = scalar_arg
self.scalar_const = scalar_const
self.scalar_index = scalar_index
def to_yaml_custom_dict(self):
if self.scalar_fn:
scalar_fn_dict = dict(kind=self.scalar_fn.kind.name.lower())
if self.scalar_fn.fn_name:
scalar_fn_dict["fn_name"] = self.scalar_fn.fn_name
if self.scalar_fn.attr_name:
scalar_fn_dict["attr_name"] = self.scalar_fn.attr_name
if self.scalar_fn.type_var:
scalar_fn_dict["type_var"] = self.scalar_fn.type_var.name
scalar_fn_dict["operands"] = list(self.scalar_fn.operands)
return dict(scalar_fn=scalar_fn_dict)
elif self.scalar_arg:
return dict(scalar_arg=self.scalar_arg.arg)
elif self.scalar_const:
return dict(scalar_const=self.scalar_const.value)
elif self.scalar_index:
return dict(scalar_index=self.scalar_index.dim)
else:
raise ValueError(f"Unexpected ScalarExpression type: {self}")
class ScalarAssign(YAMLObject):
"""An assignment to a named argument (LHS of a comprehension)."""
yaml_tag = "!ScalarAssign"
def __init__(self, arg: str, value: ScalarExpression):
self.arg = arg
self.value = value
def to_yaml_custom_dict(self):
return dict(arg=self.arg, value=self.value)
def __repr__(self):
return f"ScalarAssign({self.arg}, {self.value})"
Codebase Browser
llvm