Parsing type aliases, custom floating point format, private functions, "func.return"

mlir/astnodes.py

@@ -1,7 +1,7 @@
 """ Classes containing MLIR AST node types, fields, and conversion back to
     MLIR. """
 
-from enum import Enum, auto
+from enum import Enum
 from typing import Any, List, Union, Optional
 from lark import Token
 from lark.tree import Tree
@@ -126,12 +126,82 @@ class FloatTypeEnum(Enum):
 
 @dataclass
 class FloatType(Type):
+    pass
+
+
+@dataclass
+class StandardFloatType(FloatType):
     type: FloatTypeEnum
 
     def dump(self, indent: int = 0) -> str:
         return self.type.name
 
 
+@dataclass
+class CustomFloatType(FloatType):
+    width: int
+    exponent: int
+    mantissa: int
+    bias: int = 0
+    signed: bool = True
+    zero: bool = True
+    infinities: bool = True
+    nans: bool = True
+
+    @classmethod
+    def from_lark(cls, args: list):
+        bias: int = 0
+        signed: bool = True
+        zero: bool = True
+        infinities: bool = True
+        nans: bool = True
+        i = 0
+        while i < len(args):
+            a = args[i]
+            if isinstance(a, Token):
+                if a.value == "f":
+                    i += 1
+                    width = int(args[i])
+                elif a.value == "E":
+                    i += 1
+                    exponent = int(args[i])
+                elif a.value == "M":
+                    i += 1
+                    mantissa = int(args[i])
+                elif a.value == "B":
+                    i += 1
+                    bias = int(args[i])
+                elif a.value == "F":
+                    infinities = False
+                elif a.value == "N":
+                    nans = False
+                elif a.value == "U":
+                    signed = False
+                    zero = False
+                elif a.value == "Z":
+                    zero = True
+                else:
+                    raise ValueError(f"Unknow format specification: {a}")
+            else:
+                raise ValueError(f"Bad argument type: {a}")
+            i += 1
+        return cls(width, exponent, mantissa, bias, signed, zero, infinities, nans)
+
+
+    def dump(self, indent: int = 0) -> str:
+        suffix = ""
+        if not self.infinities:
+            suffix += "F"
+        if not self.nans:
+            suffix += "N"
+        if not self.signed:
+            suffix += "U"
+            if self.zero:
+                suffix += "Z"
+        bias = "" if self.bias == 0 else f"B{self.bias}"
+        return f"f{self.width}E{self.exponent}M{self.mantissa}{bias}{suffix}"
+
+
 @dataclass
 class TensorFloatType(Type):
     def dump(self, indent: int = 0) -> str:
@@ -185,8 +255,8 @@ def dump(self, indent: int = 0) -> str:
 
 @dataclass
 class VectorType(Type):
-    dimensions: int
-    element_type: Union[IntegerType, FloatType]
+    dimensions: List[int]
+    element_type: Union[IntegerType, FloatType, IndexType]
 
     def dump(self, indent: int = 0) -> str:
         return 'vector<%s>' % ('x'.join(
@@ -275,16 +345,6 @@ class OpaqueDialectType(Type):
     def dump(self, indent: int = 0) -> str:
         return '!%s<"%s">' % (self.dialect, self.contents)
 
-@dataclass
-class PrettyDialectType(Type):
-    dialect: str
-    type: str
-    body: List[str]
-
-    def dump(self, indent: int = 0) -> str:
-        return '!%s.%s<%s>' % (self.dialect, self.type, ', '.join(
-            dump_or_value(item, indent) for item in self.body))
-
 
 @dataclass
 class FunctionType(Type):
@@ -351,18 +411,28 @@ def dump(self, indent: int = 0) -> str:
         return '{%s}' % dump_or_value(self.value, indent)
 
 
+@dataclass
+class DenseArrayAttr(Attribute):
+    type: IntegerType | FloatType
+    value: List[bool | int | float]
+
+    def dump(self, indent: int = 0) -> str:
+        return 'array<%s: %s>' % (self.type.dump(indent),
+                                  dump_or_value(self.value, indent))
+
+
 @dataclass
 class ElementsAttr(Attribute):
     pass
 
 
 @dataclass
 class DenseElementsAttr(ElementsAttr):
-    attribute: Attribute
+    attribute: Optional[Attribute]
     type: Union[TensorType, VectorType]
 
     def dump(self, indent: int = 0) -> str:
-        return 'dense<%s> : %s' % (self.attribute.dump(indent),
+        return 'dense<%s> : %s' % (dump_or_value(self.attribute, indent),
                                    self.type.dump(indent))
 
 
@@ -619,7 +689,7 @@ class GenericModule(ModuleType):
     args: List["NamedArgument"]
     region: "Region"
     attributes: Optional[AttributeDict]
-    type: List[Type]
+    type: Type | List[Type]
     location: Optional[Location] = None
 
     def dump(self, indent=0) -> str:
@@ -634,14 +704,18 @@ def dump(self, indent=0) -> str:
         result += ')'
         if self.attributes:
             result += ' ' + dump_or_value(self.attributes, indent)
-        result += ' : ' + self.type.dump(indent)
+        if isinstance(self.type, list):
+            result += ' : ' + ', '.join(t.dump(indent) for t in self.type)
+        else:
+            result += ' : ' + self.type.dump(indent)
         if self.location:
             result += ' ' + self.location.dump(indent)
         return result
 
 
 @dataclass
 class Function(Node):
+    visibility: Optional[str]
     name: SymbolRefId
     args: Optional[List["NamedArgument"]]
     result_types: Optional[List[Type]]
@@ -651,6 +725,8 @@ class Function(Node):
 
     def dump(self, indent=0) -> str:
         result = 'func.func'
+        if self.visibility:
+            result += ' %s' % self.visibility
         result += ' %s' % self.name.dump(indent)
         arg_list = self.args if self.args else []
         result += '(%s)' % ', '.join(
@@ -999,6 +1075,8 @@ def _dump_ast_or_value(value: Any, python=True, indent: int = 0) -> str:
             '%s%s%s' %
             (_dump_ast_or_value(k, python), sep, _dump_ast_or_value(v, python))
             for k, v in value.items())
+    if value is None:
+        return ""
     return str(value)
 
 

mlir/builder/builder.py

@@ -88,9 +88,9 @@ def __init__(self):
 
     name_gen = UniqueNameGenerator(forced_prefix="_pymlir_")
 
-    F16 = mast.FloatType(type=mast.FloatTypeEnum.f16)
-    F32 = mast.FloatType(type=mast.FloatTypeEnum.f32)
-    F64 = mast.FloatType(type=mast.FloatTypeEnum.f64)
+    F16 = mast.StandardFloatType(type=mast.FloatTypeEnum.f16)
+    F32 = mast.StandardFloatType(type=mast.FloatTypeEnum.f32)
+    F64 = mast.StandardFloatType(type=mast.FloatTypeEnum.f64)
     INT32 = mast.IntegerType(32)
     INT64 = mast.IntegerType(64)
     INDEX = mast.IndexType()
@@ -143,7 +143,7 @@ def function(self, name: Optional[str] = None) -> mast.Function:
         if name is None:
             name = self.name_gen("fn")
 
-        op = mast.Function(mast.SymbolRefId(value=name), [], [], None,
+        op = mast.Function(None, mast.SymbolRefId(value=name), [], [], None,
                            mast.Region([]))
 
         self._insert_op_in_block([], op)

mlir/dialects/func.py

@@ -1,60 +1,62 @@
-
-import inspect
-import sys
-from typing import List, Tuple, Optional, Union
-from dataclasses import dataclass
-
-import mlir.astnodes as mast
-from mlir.dialect import Dialect, DialectOp, is_op
-
-Literal = Union[mast.StringLiteral, float, int, bool]
-SsaUse = Union[mast.SsaId, Literal]
-
-@dataclass
-class CallIndirectOperation(DialectOp):
-    func: mast.SymbolRefId
-    func_type: mast.FunctionType
-    args: Optional[List[SsaUse]] = None
-    argtypes: Optional[List[mast.Type]] = None
-    _syntax_ = ['func.call_indirect {func.symbol_ref_id} () : {func_type.function_type}',
-                'func.call_indirect {func.symbol_ref_id} ( {args.ssa_use_list} ) : {func_type.function_type}']
-
-
-@dataclass
-class CallOperation(DialectOp):
-    func: mast.SymbolRefId
-    func_type: mast.FunctionType
-    args: Optional[List[SsaUse]] = None
-    argtypes: Optional[List[mast.Type]] = None
-    _syntax_ = ['func.call {func.symbol_ref_id} () : {func_type.function_type}',
-                'func.call {func.symbol_ref_id} ( {args.ssa_use_list} ) : {func_type.function_type}']
-
-@dataclass
-class ConstantOperation(DialectOp):
-    value: mast.SymbolRefId
-    val_type: mast.Type
-    _syntax_ = ['func.constant {value.symbol_ref_id} : {val_typw.type}']
-
-# Note: The 'func.func' operation is defined as 'function' in mlir.lark.
-
-@dataclass
-class ReturnOperation(DialectOp):
-    values: Optional[List[SsaUse]] = None
-    types: Optional[List[mast.Type]] = None
-    _syntax_ = ['return',
-                'return {values.ssa_use_list} : {types.type_list_no_parens}']
-
-    def dump(self, indent: int = 0) -> str:
-        output = 'return'
-        if self.values:
-            output += ' ' + ', '.join([v.dump(indent) for v in self.values])
-        if self.types:
-            output += ' : ' + ', '.join([t.dump(indent) for t in self.types])
-
-        return output
-
-
-
-# Inspect current module to get all classes defined above
-func = Dialect('func', ops=[m[1] for m in inspect.getmembers(
-    sys.modules[__name__], lambda obj: is_op(obj, __name__))])
+
+import inspect
+import sys
+from typing import List, Optional, Union
+from dataclasses import dataclass
+
+import mlir.astnodes as mast
+from mlir.dialect import Dialect, DialectOp, is_op
+
+Literal = Union[mast.StringLiteral, float, int, bool]
+SsaUse = Union[mast.SsaId, Literal]
+
+@dataclass
+class CallIndirectOperation(DialectOp):
+    func: mast.SymbolRefId
+    func_type: mast.FunctionType
+    args: Optional[List[SsaUse]] = None
+    _syntax_ = ['func.call_indirect {func.symbol_ref_id} () : {func_type.function_type}',
+                'func.call_indirect {func.symbol_ref_id} ( {args.ssa_use_list} ) : {func_type.function_type}']
+
+
+@dataclass
+class CallOperation(DialectOp):
+    func: mast.SymbolRefId
+    func_type: mast.FunctionType
+    args: Optional[List[SsaUse]] = None
+    _syntax_ = ['call {func.symbol_ref_id} () : {func_type.function_type}',
+                'func.call {func.symbol_ref_id} () : {func_type.function_type}',
+                'call {func.symbol_ref_id} ( {args.ssa_use_list} ) : {func_type.function_type}',
+                'func.call {func.symbol_ref_id} ( {args.ssa_use_list} ) : {func_type.function_type}']
+
+@dataclass
+class ConstantOperation(DialectOp):
+    value: mast.SymbolRefId
+    val_type: mast.Type
+    _syntax_ = ['func.constant {value.symbol_ref_id} : {val_typw.type}']
+
+# Note: The 'func.func' operation is defined as 'function' in mlir.lark.
+
+@dataclass
+class ReturnOperation(DialectOp):
+    values: Optional[List[SsaUse]] = None
+    types: Optional[List[mast.Type]] = None
+    _syntax_ = ['return',
+                'func.return',
+                'return {values.ssa_use_list} : {types.type_list_no_parens}',
+                'func.return {values.ssa_use_list} : {types.type_list_no_parens}']
+
+    def dump(self, indent: int = 0) -> str:
+        output = 'return'
+        if self.values:
+            output += ' ' + ', '.join([v.dump(indent) for v in self.values])
+        if self.types:
+            output += ' : ' + ', '.join([t.dump(indent) for t in self.types])
+
+        return output
+
+
+
+# Inspect current module to get all classes defined above
+func = Dialect('func', ops=[m[1] for m in inspect.getmembers(
+    sys.modules[__name__], lambda obj: is_op(obj, __name__))])