mlir/opdsl: fix and inline type annotations, trim includes

[artagnon]

Summary

This patch cleans up simple issues in the opdsl/lang Python files.

Signed-off-by: Ramkumar Ramachandra <[e-mail]>

[gysit]

Looks like a good cleanup! I am not really the owner of OpDSL anymore and suggest to wait for an accept from @nicolasvasilache.

[artagnon]

Ping. This should be quite straightforward to review.

[jpienaar]

Have you tested this with Python 3.6? (I have been delayed in discussion to bump that requirement and ran into failure with type annotations on my previous python commit)

[artagnon]

I've tested it with Python 3.10, but shouldn't the CI fail my commit if it doesn't build on the minimum required version of Python?

[artagnon]

I see the following in LLVM's CMakeLists.txt:

if(LLVM_INCLUDE_TESTS)
  # Lit test suite requires at least python 3.6
  set(LLVM_MINIMUM_PYTHON_VERSION 3.6)
else()
  # FIXME: it is unknown if this is the actual minimum bound
  set(LLVM_MINIMUM_PYTHON_VERSION 3.0)
endif()

Should I remove the conditional and set the minimum required version to 3.6 uniformly, since the CI doesn't test the else codepath?

Python 3.6 isn't even available via homebrew, and I don't want to have to compile by hand: not sure how to test with 3.6.

[jpienaar]

No unfortunately this only failed for me during post submit and buildbots. I ran into same issue and had to resort to using podman (with mlir-nvidia docker image that is checked in). I should have RFC out Monday or so about bumping to 3.8 (the change in 3.7 wrt types was big one syntax wise AFAICS).

[artagnon]

Ping. Any progress on the Python version bump?

mlir/python/mlir/dialects/linalg/opdsl/lang/affine.py

Referencing a named symbol:
   >>> S.foobar
   Symbol(foobar)
   >>> S.a is S.b
   False
   >>> S.a is S.a
   True
 """
 
-from typing import Callable, Dict, Optional, Tuple, Union
+from typing import Dict, Optional
 
 from ..... import ir as _ir
 
 __all__ = [
     "AffineBuildState",
     "AffineExprDef",
     "D",
     "DimDef",
class AffineBuildState:
   AffineBuildState. In that case, any affine expressions built will inherit
   symbol and dim bindings from the global state and will update both as new
   ones are discovered. This allows for building expressions across contexts
   which share a common symbol and dim space.
   """
 
   def __init__(self,
                *,
-               global_state: "AffineBuildState" = None,
+               global_state: Optional["AffineBuildState"] = None,
                allow_new_symbols: bool = True,
                allow_new_dims: bool = True):
     if not global_state:
-      self.all_symbols = dict()  # type: Dict[str, int]
-      self.all_dims = dict()  # type: Dict[str, int]
+      self.all_symbols: Dict[str, int] = dict()
+      self.all_dims: Dict[str, int] = dict()
     else:
       # Alias the global dict.
       self.all_symbols = global_state.all_symbols
       self.all_dims = global_state.all_dims
 
     # Map of symbols and dims in the current build.
-    self.local_symbols = dict()  # type: Dict[str, int]
-    self.local_dims = dict()  # type: Dict[str, int]
+    self.local_symbols: Dict[str, int] = dict()
+    self.local_dims: Dict[str, int] = dict()
     self.allow_new_symbols = allow_new_symbols
     self.allow_new_dims = allow_new_dims
 
   def get_dim(self, dimname: str) -> int:
     """Gets the dim position given a name."""
     pos = self.all_dims.get(dimname)
     if pos is None:
       if not self.allow_new_dims:
class AffineBinaryExprDef(AffineExprDef):
   def __repr__(self):
     return f"{self.ir_ctor.__name__}({repr(self.lhs)}, {repr(self.rhs)})"
 
 
 class DimDef(AffineExprDef):
   """Represents a named dimension.
 
   """
-  ALL_DIMS = dict()  # type: Dict[str, "DimDef"]
+  ALL_DIMS: Dict[str, "DimDef"] = dict()
 
   def __new__(cls, dimname: str):
     existing = cls.ALL_DIMS.get(dimname)
     if existing is not None:
       return existing
     new = super().__new__(cls)
     new.dimname = dimname
     cls.ALL_DIMS[dimname] = new
"""Represents a named symbol.
     >>> s1
     Symbol(s1)
     >>> s2 = SymbolDef("s2")
     >>> s1 is s2
     False
     >>> s1 is SymbolDef("s1")
     True
   """
-  ALL_SYMBOLS = dict()  # type: Dict[str, "SymbolDef"]
+  ALL_SYMBOLS: Dict[str, "SymbolDef"] = dict()
 
   def __new__(cls, symname: str):
     existing = cls.ALL_SYMBOLS.get(symname)
     if existing is not None:
       return existing
     new = super().__new__(cls)
     new.symname = symname
     cls.ALL_SYMBOLS[symname] = new

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py

class TensorReduceFn(TensorExpression):
   """Application of a reduction function.
 
   This captures the lhs (initial value) separately from the rhs.
   """
 
   def __init__(self, reduce_use: "ReduceFnUse",
                args: Sequence[TensorExpression]):
     self.reduce_use = reduce_use
-    self.lhs = None  # type: Optional[TensorUse]
+    self.lhs: Optional[TensorUse] = None
     self.args = args
 
   def to_scalar_expression(self) -> ScalarExpression:
     if self.lhs is None:
       raise ValueError(f"Cannot scalarize a TensorReduceFn that has not been "
                        f"bound to its lhs: {self}")
     full_args = [self.lhs.to_scalar_expression()
                 ] + [arg.to_scalar_expression() for arg in self.args]
def __init__(self,
                type_var: Optional[TypeVar] = None,
                size_exprs: Optional[Sequence[AffineExprDef]] = None,
                index_dims: Optional[Sequence[DimDef]] = None,
                default_indices: Optional[Sequence[int]] = None,
                default_fn: Optional[str] = None):
     if type_var and not isinstance(type_var, TypeVar):
       raise ValueError(
           f"OperandDef requires a TypeVar but got {repr(type_var)}")
-    self.owner = None  # type: Optional["LinalgOpDef"]
+    self.owner: Optional["LinalgOpDef"] = None
     self.type_var = type_var
     self.size_exprs = size_exprs
     self.index_dims = index_dims
     self.default_indices = default_indices
     self.default_fn = default_fn
     self.kind = kind
-    self.name = None  # type: Optional[str]
-    self.registered_index = -1  # type: int
+    self.name: Optional[str] = None
+    self.registered_index: int = -1
 
   def attach(self, index: int, name: str, owner: "LinalgOpDef"):
     if self.owner:
       raise ValueError(f"OperandDef already registered with an op: {self}")
     self.registered_index = index
     self.name = name
     self.owner = owner
 
def __getitem__(self, dims: Sequence[AffineExprDef]) -> TensorUse:
     return TensorUse(self.operand_def, exprs)
 
   def __setitem__(self, dims: Sequence[AffineExprDef], value: TensorExpression):
     """Creates a new 1:1 comprehension by binding this tensor to an expression.
 
     Note that due to the way assignment works in Python, we have to capture
     direct assignment as a setitem on the TensorDef.
     """
+    assert self.operand_def.owner, "TensorDef is not registered with an op"
     if not isinstance(value, TensorExpression):
       raise ValueError(f"Only TensorExpressions can be assigned to TensorDefs. "
                        f"Got: {repr(value)}")
-    use = self[dims]
-    comp = Comprehension((use, value))
+    comp = Comprehension((self[dims], value))
     self.operand_def.owner.comprehensions.append(comp)
 
 
 class ScalarDef(TensorExpression):
   """Scalar operand definition.
 
   Scalar operands are forwarded to the body of the structured op as they are.
   A unique name identifies the scalars and an index determines their position in

 # Operation definition.
 ###############################################################################
 
 
 class Comprehension:
   """Represents a single comprehension."""
 
   def __init__(self, *bindings: Tuple[TensorUse, TensorExpression]):
-    self.definitions = list()  # List[TensorUse]
-    self.values = list()  # List[TensorExpression]
+    self.definitions: List[TensorUse] = list()
+    self.values: List[TensorExpression] = list()
 
     # Find the lhs to reduction rhs.
     for assign, value in bindings:
       if isinstance(value, TensorReduceFn):
         if value.lhs:
           raise ValueError(f"Reduction expression already assigns: {value}")
         value.lhs = assign
       self.definitions.append(assign)
       self.values.append(value)
 
   @property
   def all_reduction_dims(self) -> Set[Tuple[DimDef, ...]]:
     """Gets the reduction dims for the comprehension or None."""
-    result = set()
+    result: Set[Tuple[DimDef, ...]] = set()
     for use in self.values:
       if isinstance(use, TensorReduceFn):
         result.add(use.reduce_use.reduce_dims)
       else:
         result.add(tuple())
     return result
 
   def __repr__(self):

 
 class OpMetadataDef(YAMLObject):
   """Metadata about the op (generally not behavior impacting)."""
   yaml_tag = "!LinalgOpMetadata"
 
   def __init__(self, name: str, cpp_class_name: Optional[str],
                doc: Optional[str]):
     self.name = name
-    self.cpp_class_name = cpp_class_name if cpp_class_name is not None else name
+    self.cpp_class_name = cpp_class_name if cpp_class_name else name
     self.doc = doc
-    self.implements = []  # type: List[OpInterfaceDef]
-    self.defines = []  # type: List[OpDefinitionsDef]
+    self.implements: List[OpInterfaceDef] = []
+    self.defines: List[OpDefinitionDef] = []
 
   def to_yaml_custom_dict(self):
     d = dict(
         name=self.name,
         cpp_class_name=self.cpp_class_name,
         doc=self.doc,
     )
     if self.implements:
       d["implements"] = [intr.cpp_name for intr in self.implements]
     if self.defines:
       d["defines"] = [defi.def_name for defi in self.defines]
     return d
 
 
 class LinalgOpDef:
   """Definition of a linalg op."""
 
   def __init__(self,
                name: str,
                cpp_class_name: Optional[str] = None,
                doc: Optional[str] = None):
     self.metadata = OpMetadataDef(
         name=name, cpp_class_name=cpp_class_name, doc=doc)
-    self.registered_operands = dict()  # type: Dict[str, OperandDef]
-    self.domain = list()  # type: List[DimDef]
-    self.comprehensions = list()  # type: List[Comprehension]
+    self.registered_operands: Dict[str, OperandDef] = dict()
+    self.domain: List[DimDef] = list()
+    self.comprehensions: List[Comprehension] = list()
     self._affine_state = AffineBuildState()
 
   def add_operand(self, name: str, operand: OperandDef):
     """Registers an operand."""
     if name in self.registered_operands:
       raise ValueError(f"The operand {name} is already registered "
                        f"to {self.registered_operands['name']}")
     structured_op_methods = [

mlir/python/mlir/dialects/linalg/opdsl/lang/config.py

class OperandDefConfig(YAMLObject):
   """Wrapper containing an operand definition with additional state."""
   yaml_tag = "!LinalgOperandDefConfig"
 
   def __init__(self,
                operand_def: OperandDef,
                shape_map: Optional[_ir.AffineMap] = None,
                index_attr_map: Optional[_ir.AffineMap] = None):
     self.operand_def = operand_def
-    self.shape_map = shape_map  # type: Optional[_ir.AffineMap]
-    self.index_attr_map = index_attr_map  # type: Optional[_ir.AffineMap]
-    self.indexing_map = None  # type: Optional[_ir.AffineMap]
+    self.shape_map: Optional[_ir.AffineMap] = shape_map
+    self.index_attr_map: Optional[_ir.AffineMap] = index_attr_map
+    self.indexing_map: Optional[_ir.AffineMap] = None
 
   @property
   def name(self) -> str:
     return self.operand_def.name
 
   @property
   def kind(self) -> OperandKind:
     return self.operand_def.kind
class LinalgStructuredOpConfig(YAMLObject):
 
   def __init__(self,
                comprehension: Comprehension,
                domain: Sequence[DimDef],
                registered_operands: Sequence[OperandDef],
                context: Optional[_ir.Context] = None):
     self.context = context if context is not None else _ir.Context()
     self.affine_state = AffineBuildState()
-    self.writes = list()  # type: List[Tuple[TensorUse, TensorExpression]]
-    self.operands = dict()  # type: Dict[OperandDef, OperandDefConfig]
-    self.uses = dict()  # type: Dict[TensorUse, TensorUseConfig]
+    self.writes: List[Tuple[TensorUse, TensorExpression]] = list()
+    self.operands: Dict[OperandDef, OperandDefConfig] = dict()
+    self.uses: Dict[TensorUse, TensorUseConfig] = dict()
 
     # Compute the ordered set of writes and collect the tensor, capture, dims,
     # and index uses.
     collected_tensor_uses = set()
     collected_scalar_uses = set()
     collected_dim_uses = set()
     collected_indices = set()
     for write_use, read_use in zip(comprehension.definitions,
for write_use, _ in self.writes:
       if write_tensor_config.indexing_map:
         raise ValueError(
             f"Unexpected multi-write to a single tensor: {write_tensor_config}")
       write_tensor_config.indexing_map = self.uses[write_use].indexing_map
 
     # For each read use, propagate the indexing maps from the use to the
     # tensor, ensuring that there are not conflicts.
     for _, read_expr in self.writes:
-      read_uses = set()  # type: Set[TensorUse]
+      read_uses: Set[TensorUse] = set()
       read_expr.collect_tensor_uses(read_uses)
       for read_use in read_uses:
         read_operand_config = self.operands[read_use.operand_def]
         if (read_operand_config.indexing_map and
             read_operand_config.indexing_map !=
             self.uses[read_use].indexing_map):
           raise ValueError(
               f"Unexpected multi-read of a tensor with different accesses:"

mlir/python/mlir/dialects/linalg/opdsl/lang/dsl.py

 #  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
 
-from typing import Dict, List, Sequence, Union
+from typing import Sequence, Union
 
 from contextlib import contextmanager
 import functools
 import inspect
 import threading
 
 from ..... import ir
 from ...._ods_common import get_op_result_or_value as _get_op_result_or_value, get_op_results_or_values as _get_op_results_or_values

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

 #  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
 
 from typing import Callable, Dict, List, Sequence, Tuple, Union
 
 from .....ir import *
 
 from .... import func
 from .... import linalg
 from .... import math
 from .... import arith
 from .... import complex
-from ...._ods_common import get_op_result_or_value as _get_op_result_or_value, get_op_results_or_values as _get_op_results_or_values
 
 from .scalar_expr import *
 from .config import *
 from .comprehension import *
 import numpy as np
 
 __all__ = [
     "emit_generic_structured_op",
def prepare_common_structured_op(op_config: LinalgStructuredOpConfig,
   if len(ins) != len(in_arg_defs):
     raise ValueError(f"Expected {len(in_arg_defs)} inputs but got "
                      f"{len(ins)} for {op_config}")
   if outs and len(outs) != len(out_arg_defs):
     raise ValueError(f"Expected {len(out_arg_defs)} outputs but got "
                      f"{len(outs)} for {op_config}")
 
   # Compute a replacement list for all index attribute symbols.
-  expressions = []  # type: Sequence[AffineExpr]
-  replacements = []  # type: Sequence[AffineExpr]
+  expressions: Sequence[AffineExpr] = []
+  replacements: Sequence[AffineExpr] = []
   for index_attr in index_attr_arg_defs:
     index_attr_vals = index_attr.operand_def.default_indices
     if index_attr.name in attrs:
       index_attr_vals = attrs.get(index_attr.name)
     assert index_attr_vals, "Index attribute has no value"
     if not all(isinstance(value, int) for value in index_attr_vals):
       raise ValueError(f"Attribute {index_attr.name} needs to be of type "
                        f"Sequence[int] but got {type(index_attr_vals)}")
-    results = index_attr.index_attr_map.results  # type: AffineExprList
+    results: AffineExprList = index_attr.index_attr_map.results
     if len(index_attr_vals) != len(results):
       raise ValueError(f"Attribute {index_attr.name} has length {len(results)} "
                        f"but got {len(index_attr_vals)} values")
     for expr, value in zip(results, index_attr_vals):
       expressions.append(expr)
       replacements.append(AffineConstantExpr.get(value))
 
   # Replace all index attribute symbols by their value.
   # TODO: Add support for shape symbols.
-  indexing_maps = []  # type: Sequence[AffineMap]
+  indexing_maps: Sequence[AffineMap] = []
   for curr in op_config.indexing_maps:
     for expression, replacement in zip(expressions, replacements):
       curr = curr.replace(expression, replacement, curr.n_dims, curr.n_symbols)
     indexing_maps.append(curr)
 
   # TODO: Linalg verification does not currently allow symbols.
   # Compress them for now and verify none are left.
   indexing_maps = AffineMap.compress_unused_symbols(indexing_maps,
                                                     Context.current)
   if any(indexing_map.n_symbols != 0 for indexing_map in indexing_maps):
     raise ValueError(f"Expected indexing_maps to use no symbols after "
                      f"replacement and compression but got {indexing_maps}")
 
   outs, out_types = _infer_structured_outs(op_config, in_arg_defs, ins,
                                            out_arg_defs, outs)
 
   result_types = [t for t in out_types if isa(RankedTensorType, t)]
 
   # Initialize the type dictionary with the predefined types.
-  type_mapping = dict()  # type: Dict[str, Type]
+  type_mapping: Dict[str, Type] = dict()
   type_mapping["F32"] = F32Type.get()
   type_mapping["F64"] = F64Type.get()
   type_mapping["I32"] = IntegerType.get_signless(32)
   type_mapping["I64"] = IntegerType.get_signless(64)
 
   # Extract type vars for input/output based types.
-  block_arg_types = list()  # type: List[Type]
+  block_arg_types: List[Type] = list()
   for arg_def, arg_element_type in zip(in_arg_defs + out_arg_defs,
                                        _get_types_from_values(*ins, *outs)):
     _add_type_mapping(arg_def, arg_element_type, type_mapping, block_arg_types)
 
   # Emit the generic op.
   # TODO: Support emission of pure memref form.
   indexing_maps_attr = ArrayAttr.get(
       [AffineMapAttr.get(am) for am in indexing_maps])
   iterator_types_attr = ArrayAttr.get([
       Attribute.parse(f"#linalg.iterator_type<{s}>")
       for s in op_config.iterator_types
   ])
 
   # Compute the index attributes used when emitting a named structured op.
-  index_attrs = {}  # type: Dict[str, DenseElementAttr]
+  index_attrs: Dict[str, DenseElementsAttr] = {}
   for index_attr in index_attr_arg_defs:
     index_attr_vals = attrs.get(index_attr.name)
     # Only forward attributes set to a non-default value.
     if index_attr_vals:
       array = np.array(index_attr_vals, dtype=np.int64)
       index_attrs[index_attr.name] = DenseElementsAttr.get(array)
 
   # Compute the function attribute mapping.
class _BodyBuilder:
   """Constructs a structured op body by evaluating assignments."""
 
   def __init__(self, type_mapping: Dict[str, Type],
                block_arg_mapping: Dict[str, Value], fn_attr_mapping: Dict[str,
                                                                           str]):
     self.type_mapping = type_mapping
     self.block_arg_mapping = block_arg_mapping
     self.fn_attr_mapping = fn_attr_mapping
-    self.yield_mapping = dict()  # type: Dict[str, Value]
+    self.yield_mapping: Dict[str, Value] = dict()
 
   def assign(self, assignment: ScalarAssign):
     if assignment.arg in self.yield_mapping:
       raise ValueError(
           f"Multiple assignments to the same argument are forbidden: "
           f"{assignment}")
     self.yield_mapping[assignment.arg] = self.expression(assignment.value)
 

mlir/python/mlir/dialects/linalg/opdsl/lang/types.py

 ]
 
 
 class TypeVar:
   """A replaceable type variable.
 
   Type variables are uniqued by name.
   """
-  ALL_TYPEVARS = dict()  # type: Dict[str, "TypeVar"]
+  ALL_TYPEVARS: Dict[str, "TypeVar"] = dict()
 
   def __new__(cls, name: str):
     existing = cls.ALL_TYPEVARS.get(name)
     if existing is not None:
       return existing
     new = super().__new__(cls)
     new.name = name
     cls.ALL_TYPEVARS[name] = new