[mlir][linalg][python] Add attribute support to the OpDSL.

Extend the OpDSL with index attributes. After tensors and scalars, index attributes are the third operand type. An index attribute represents a compile-time constant that is limited to index expressions. A use cases are the strides and dilations defined by convolution and pooling operations.

The patch only updates the OpDSL. The C++ yaml codegen is updated by a followup patch.

Differential Revision: https://reviews.llvm.org/D104711

Author: Tobias Gysi

mlir/include/mlir-c/AffineMap.h

@@ -169,6 +169,13 @@ mlirAffineMapGetMajorSubMap(MlirAffineMap affineMap, intptr_t numResults);
 MLIR_CAPI_EXPORTED MlirAffineMap
 mlirAffineMapGetMinorSubMap(MlirAffineMap affineMap, intptr_t numResults);
 
+/// Apply AffineExpr::replace(`map`) to each of the results and return a new
+/// new AffineMap with the new results and the specified number of dims and
+/// symbols.
+MLIR_CAPI_EXPORTED MlirAffineMap mlirAffineMapReplace(
+    MlirAffineMap affineMap, MlirAffineExpr expression,
+    MlirAffineExpr replacement, intptr_t numResultDims, intptr_t numResultSyms);
+
 /// Returns the simplified affine map resulting from dropping the symbols that
 /// do not appear in any of the individual maps in `affineMaps`.
 /// Asserts that all maps in `affineMaps` are normalized to the same number of

mlir/lib/Bindings/Python/IRAffine.cpp

@@ -654,6 +654,14 @@ void mlir::python::populateIRAffine(py::module &m) {
                  mlirAffineMapGetMinorSubMap(self, nResults);
              return PyAffineMap(self.getContext(), affineMap);
            })
+      .def("replace",
+           [](PyAffineMap &self, PyAffineExpr &expression,
+              PyAffineExpr &replacement, intptr_t numResultDims,
+              intptr_t numResultSyms) {
+             MlirAffineMap affineMap = mlirAffineMapReplace(
+                 self, expression, replacement, numResultDims, numResultSyms);
+             return PyAffineMap(self.getContext(), affineMap);
+           })
       .def_property_readonly(
           "is_permutation",
           [](PyAffineMap &self) { return mlirAffineMapIsPermutation(self); })

mlir/lib/CAPI/IR/AffineMap.cpp

@@ -138,6 +138,15 @@ MlirAffineMap mlirAffineMapGetMinorSubMap(MlirAffineMap affineMap,
   return wrap(unwrap(affineMap).getMinorSubMap(numResults));
 }
 
+MlirAffineMap mlirAffineMapReplace(MlirAffineMap affineMap,
+                                   MlirAffineExpr expression,
+                                   MlirAffineExpr replacement,
+                                   intptr_t numResultDims,
+                                   intptr_t numResultSyms) {
+  return wrap(unwrap(affineMap).replace(unwrap(expression), unwrap(replacement),
+                                        numResultDims, numResultSyms));
+}
+
 void mlirAffineMapCompressUnusedSymbols(
     MlirAffineMap *affineMaps, intptr_t size, void *result,
     void (*populateResult)(void *res, intptr_t idx, MlirAffineMap m)) {

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

@@ -9,6 +9,7 @@
 """
 
 from typing import Any, Dict, List, Optional, Sequence, Set, Tuple
+from enum import Enum
 
 from mlir import ir as _ir
 
@@ -133,18 +134,31 @@ def __repr__(self):
     return f"{self.tensor_name}[{', '.join([repr(i) for i in self.indices])}]"
 
 
+class OperandKind(Enum):
+  InputTensor = 0
+  Scalar = 1
+  OutputTensor = 2
+  Attribute = 3
+
+
 class OperandDef:
-  """Definition of a Tensor or Scalar operand passed to an operation."""
+  """Definition of an operand passed to an operation.
+
+  Keep the meta information of Tensor, Scalar, and Attribute operands and
+  provide the shared registration functionality.
+  """
 
-  def __init__(self, type_var: TypeVar, shape: Sequence[AffineExprDef],
-               scalar: bool, output: bool):
+  def __init__(self,
+               kind: OperandKind,
+               type_var: TypeVar,
+               size_exprs: Optional[Sequence[AffineExprDef]] = None):
     if not isinstance(type_var, TypeVar):
-      raise ValueError(f"OperandDef requires a TypeVar. Got: {repr(type_var)}")
+      raise ValueError(
+          f"OperandDef requires a TypeVar but got {repr(type_var)}")
     self.owner = None  # type: Optional["LinalgOpDef"]
     self.type_var = type_var
-    self.shape = shape
-    self.scalar = scalar
-    self.output = output
+    self.size_exprs = size_exprs
+    self.kind = kind
     self.name = None  # type: Optional[str]
     self.registered_index = -1  # type: int
 
@@ -159,25 +173,26 @@ def __hash__(self):
     return hash(id(self))
 
   def __repr__(self):
-    output = "OUTPUT " if self.output else ""
-    scalar = "SCALAR " if self.scalar else ""
-    return (f"{self.name}:OperandDef({output}{scalar}"
-            f"{repr(self.type_var)}, shape={self.shape})")
+    return (f"{self.name}:OperandDef(kind={self.kind.name}, "
+            f"type={repr(self.type_var)}, size_exprs={self.size_exprs})")
 
 
 class TensorDef:
   """Tensor operand definition.
 
   Tensor operands are indexed using the associated indexing_map when forwarded
   to the body of the structured op. A unique name identifies the tensor operands
-  and an index determines their position in the operation's parameter list.
+  and an index determines their position in the operation's parameter list. A
+  tensor definition takes type, a shape, and an optional flag to mark output
+  tensors.
   """
 
   def __init__(self,
                type_var: TypeVar,
                *shape: AffineExprDef,
                output: bool = False):
-    self.operand_def = OperandDef(type_var, shape, False, output)
+    kind = OperandKind.OutputTensor if output else OperandKind.InputTensor
+    self.operand_def = OperandDef(kind, type_var, size_exprs=shape)
 
   def __getitem__(self, dims) -> TensorUse:
     assert self.operand_def.owner, "TensorDef is not attached to an op"
@@ -221,7 +236,7 @@ class ScalarDef(TensorExpression):
   """
 
   def __init__(self, type_var: TypeVar):
-    self.operand_def = OperandDef(type_var, (), True, False)
+    self.operand_def = OperandDef(OperandKind.Scalar, type_var)
 
   @property
   def scalar_name(self) -> str:
@@ -233,6 +248,22 @@ def to_scalar_expression(self) -> ScalarExpression:
     return ScalarArg(self.scalar_name).expr()
 
 
+class AttributeDef:
+  """Index Attribute definition.
+
+  Index attributes provide a way to define and set symbols that can be used in
+  indexing expressions. Every attribute specifies a tuple of symbols that at
+  compile-time are replaced by integer values.
+  """
+  yaml_tag = "!LinalgAttributeDef"
+
+  def __init__(self, *sizes: SymbolDef):
+    if any(not isinstance(size, SymbolDef) for size in sizes):
+      raise ValueError(f"AttributeDef requires sizes of type SymbolDef but got "
+                       f"{type(sizes)}")
+    self.operand_def = OperandDef(OperandKind.Attribute, I64, size_exprs=sizes)
+
+
 class Comprehension:
   """Represents a single comprehension."""
 
@@ -303,7 +334,7 @@ class ReduceFnType:
   def __init__(self, operator: PrimFnType, *reduce_dims: DimDef):
     """Initializes the ReduceFn with a primitive function and dims."""
     if not isinstance(operator, PrimFnType):
-      raise ValueError(f"Reduce expected a Prim operator. Got: {operator}")
+      raise ValueError(f"Reduce expected a Prim operator but got {operator}")
     self.operator = operator
     self.reduce_dims = tuple(reduce_dims)
 
@@ -353,7 +384,7 @@ def __init__(self, value: Any):
         self.value = str(
             _ir.IntegerAttr.get(_ir.IntegerType.get_signless(64), int(value)))
       else:
-        raise ValueError(f"const requires int or float. Got: {type(value)}")
+        raise ValueError(f"const requires int or float but got {type(value)}")
 
   def to_scalar_expression(self) -> ScalarExpression:
     return ScalarConst(self.value).expr()
@@ -475,21 +506,22 @@ def __init__(self,
     self.comprehensions = list()  # type: List[Comprehension]
     self._affine_state = AffineBuildState()
 
-  @property
-  def outputs(self) -> Sequence[OperandDef]:
-    return [
-        operand for operand in self.registered_operands.values()
-        if operand.output
-    ]
-
   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']}")
-    if not operand.output and self.outputs:
-      raise ValueError(f"The operand {name} is an input registered after "
-                       f"the output {self.outputs[-1]}")
+    # Ensure output tensors are registered after input tensors and scalars and
+    # attributes are registered after all other operand types.
+    registered_kinds = [
+        operand.kind.value for operand in self.registered_operands.values()
+    ]
+    if registered_kinds:
+      maximum = max(registered_kinds)
+      if maximum > operand.kind.value and maximum > OperandKind.Scalar.value:
+        raise ValueError(
+            f"The operand {name} of kind {operand.kind.name} is registered "
+            f"after an operand of kind {OperandKind(maximum).name}")
     operand.attach(len(self.registered_operands), name, self)
     self.registered_operands[name] = operand