[mlir][OpDSL] Separate `ReduceFn` and `ReduceFnUse`.

[gysit]

Summary

The revision distinguishes ReduceFn and ReduceFnUse. The latter has the reduction dimensions attached while the former specifies the arithmetic function only. This separation allows us to adapt the reduction syntax a little bit and specify the reduction dimensions using square brackets (in contrast to the round brackets used for the values to reduce). It als is a preparation to add reduction function attributes to OpDSL. A reduction function attribute shall only specify the arithmetic function and not the reduction dimensions.

Example:

ReduceFn.max_unsigned(D.kh, D.kw)(...)

changes to:

ReduceFn.max_unsigned[D.kh, D.kw](...)

Depends On D115240

[gysit]

Remove debug print statement.

[gysit]

Improve doc.

[stellaraccident]

Nice improvement - thanks! (and sorry for review latency)

mlir/docs/Dialects/Linalg/OpDSL.md

 *   `ArithFn.max_unsigned(a, b)`
 *   `ArithFn.min_unsigned(a, b)`
 
 Reduction functions can appear as the outer-most function on the RHS:
 
 *   `ReduceFn.add` (also overloading the inplace `+=` on a LHS)
 *   `ReduceFn.mul`
 *   `ReduceFn.max`
+*   `ReduceFn.min`
+*   `ReduceFn.max_unsigned`
+*   `ReduceFn.min_unsigned`
 
 Type functions cast the `operand` to the target type `TypeVar` applying signed
 or unsigned semantics for integer operands:
 
 *   `TypeFn.cast(TypeVar, operand)`
 *   `TypeFn.cast_unsigned(TypeVar, operand)`
 
 There are also special forms:

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

def collect_dim_uses(self, uses: Set["DimDef"]):
     def visit_dim_def(dim_def):
       if isinstance(dim_def, DimDef):
         uses.add(dim_def)
 
     def visit_affine_exprs(expr):
       if isinstance(expr, TensorUse):
         for ind in expr.indices:
           ind.visit_affine_exprs(visit_dim_def)
-      if isinstance(expr, ReduceApply):
-        for ind in expr.reduce.reduce_dims:
+      if isinstance(expr, TensorReduceFn):
+        for ind in expr.reduce_fn.reduce_dims:
           ind.visit_affine_exprs(visit_dim_def)
 
     self.visit_tensor_exprs(visit_affine_exprs)
 
   def collect_tensor_uses(self, uses: Set["TensorUse"]):
     """Collects all TensorUses reachable through this expression."""
 
     def visit_tensor_use(expr):
def to_scalar_expression(self) -> ScalarExpression:
     return ScalarArg(self.tensor_name).expr()
 
   @property
   def tensor_name(self) -> str:
     name = self.operand_def.name
     assert name is not None, "TensorDef not attached"
     return name
 
-  def __iadd__(self, rhs: TensorExpression) -> TensorExpression:
-    return ReduceFn.add(*self._compute_reduce_dims(rhs))(rhs)
+  def __iadd__(self, rhs: TensorExpression) -> "TensorReduceFn":
+    return ReduceFnUse(ArithFn.add, *self._compute_reduce_dims(rhs))(rhs)
 
   def _compute_reduce_dims(self, rhs: TensorExpression) -> Set[DimDef]:
     """For implicit reductions, computes default reduction dims.
 
     Assumes that the rhs is the expression being reduced and self is being
     reduced into. Any indices referenced on the rhs and not in self are
     considered reduction dims and will be ordered as encountered on the rhs.
     """
class Comprehension:
   """Represents a single comprehension."""
 
   def __init__(self, *bindings: Tuple[TensorUse, TensorExpression]):
     self.definitions = list()  # List[TensorUse]
     self.values = list()  # List[TensorExpression]
 
     # Find the lhs to reduction rhs.
     for assign, value in bindings:
-      if isinstance(value, ReduceApply):
+      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()
     for use in self.values:
-      if isinstance(use, ReduceApply):
-        result.add(use.reduce.reduce_dims)
+      if isinstance(use, TensorReduceFn):
+        result.add(use.reduce_use.reduce_dims)
       else:
         result.add(tuple())
     return result
 
   def __repr__(self):
     if len(self.definitions) > 1:
       defs_repr = f"({', '.join(repr(d) for d in self.definitions)})"
       values_repr = f"({', '.join(repr(v) for v in self.values)})"
class ArithFnType:
   """Arithmetic operations."""
 
   def __init__(self, fn_name: str):
     self.fn_name = fn_name
 
   def __call__(self, *args) -> "TensorArithFn":
     return TensorArithFn(self, args)
 
-  def reduce(self, *reduce_dims: DimDef):
-    """Shortcut to create a Reduce operation from this function."""
-    return ReduceFnType(self, *reduce_dims)
-
   def __repr__(self):
     return f"{self.fn_name}"
 
 
 class ArithFn:
   add = ArithFnType("add")
   exp = ArithFnType("exp")
   log = ArithFnType("log")
   mul = ArithFnType("mul")
   max = ArithFnType("max")
   min = ArithFnType("min")
   sub = ArithFnType("sub")
   max_unsigned = ArithFnType("max_unsigned")
   min_unsigned = ArithFnType("min_unsigned")
 
 
-class ReduceFnType:
-  """A reduction operator that reduces into its LHS from its RHS."""
+class ReduceFnUse:
+  """A reduction function assigned to reduction dimensions."""
 
-  def __init__(self, operator: ArithFnType, *reduce_dims: DimDef):
-    """Initializes the ReduceFn with an airthmetic function and dims."""
-    if not isinstance(operator, ArithFnType):
-      raise ValueError(f"Reduce expected a ArithFnType but got {operator}")
-    self.operator = operator
-    self.reduce_dims = tuple(reduce_dims)
+  def __init__(self, arith_fn: ArithFnType, *reduce_dims: DimDef):
+    self.arith_fn = arith_fn
+    self.reduce_dims = reduce_dims
 
   def __call__(self, *args: TensorExpression):
-    return ReduceApply(self, args)
+    return TensorReduceFn(self, args)
 
   def __repr__(self):
-    return (f"reduce_{self.operator.fn_name}"
+    return (f"reduce_{self.arith_fn.fn_name}"
             f"({', '.join(repr(d) for d in self.reduce_dims)})")
 
 
+class ReduceFnType:
+  """Reduction operations.
+
+  All reduction operations reduces their RHS into their LHS.
+  """
+
+  def __init__(self, arith_fn: ArithFnType):
+    if not isinstance(arith_fn, ArithFnType):
+      raise ValueError(f"Reduce expected a ArithFnType but got {arith_fn}")
+    self.arith_fn = arith_fn
+
+  def __getitem__(self, reduce_dims: Tuple[DimDef]) -> ReduceFnUse:
+    return ReduceFnUse(self.arith_fn, *reduce_dims)
+
+  def __repr__(self):
+    return (f"reduce_{self.arith_fn.fn_name}")
+
+
 class ReduceFn:
-  add = ArithFn.add.reduce
-  mul = ArithFn.mul.reduce
-  max = ArithFn.max.reduce
-  min = ArithFn.min.reduce
-  max_unsigned = ArithFn.max_unsigned.reduce
-  min_unsigned = ArithFn.min_unsigned.reduce
+  add = ReduceFnType(ArithFn.add)
+  mul = ReduceFnType(ArithFn.mul)
+  max = ReduceFnType(ArithFn.max)
+  min = ReduceFnType(ArithFn.min)
+  max_unsigned = ReduceFnType(ArithFn.max_unsigned)
+  min_unsigned = ReduceFnType(ArithFn.min_unsigned)
 
 
 class TensorArithFn(TensorExpression):
   """Application of an arithmetic function."""
 
   def __init__(self, arith_fn: ArithFnType, args: Sequence[TensorExpression]):
     self.arith_fn = arith_fn
     self.args = tuple(args)
class index(TensorExpression):
   def to_scalar_expression(self) -> ScalarExpression:
     assert self.dim != -1, "Dimension name not resolved"
     return ScalarIndex(self.dim).expr()
 
   def __repr__(self):
     return f"index({repr(self.dim)})"
 
 
-class ReduceApply(TensorExpression):
-  """Application of a reduction.
+class TensorReduceFn(TensorExpression):
+  """Application of a reduction operation.
 
-  This captures the lhs separately (initial value) separately from the rhs.
+  This captures the lhs (initial value) separately from the rhs.
   """
 
-  def __init__(self, reduce: ReduceFnType, args: Sequence[TensorExpression]):
-    self.reduce = reduce
+  def __init__(self, reduce_use: ReduceFnUse, args: Sequence[TensorExpression]):
+    self.reduce_use = reduce_use
     self.lhs = None  # type: Optional[TensorUse]
     self.args = tuple(args)
 
   def to_scalar_expression(self) -> ScalarExpression:
     if self.lhs is None:
-      raise ValueError(f"Cannot scalarize a ReduceApply that has not been "
+      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]
-    return ScalarArithFn(self.reduce.operator.fn_name, *full_args).expr()
+    return ScalarArithFn(self.reduce_use.arith_fn.fn_name, *full_args).expr()
 
   def visit_tensor_exprs(self, callback):
     for arg in self.args:
       arg.visit_tensor_exprs(callback)
 
   def __repr__(self):
-    return f"{repr(self.reduce)}({', '.join(repr(a) for a in self.args)})"
+    return f"{repr(self.reduce_use)}({', '.join(repr(a) for a in self.args)})"
 
 
 class OpInterfaceDef:
   """An interface that an op implements."""
 
   def __init__(self, cpp_name: str):
     self.cpp_name = cpp_name
 

mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py

def pooling_nhwc_max(
     dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs max pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
-  O[D.n, D.oh, D.ow, D.c] = ReduceFn.max(D.kh, D.kw)(
+  O[D.n, D.oh, D.ow, D.c] = ReduceFn.max[D.kh, D.kw](
       TypeFn.cast(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_nhwc_max_unsigned(
     I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
                 S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
     strides=IndexAttrDef(S.SH, S.SW),
     dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs unsigned max pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
-  O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_unsigned(D.kh, D.kw)(
+  O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_unsigned[D.kh, D.kw](
       TypeFn.cast_unsigned(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_nchw_max(
     I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH,
                 S.OW * S.SW + S.KW * S.DW),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.C, S.OH, S.OW, output=True),
     strides=IndexAttrDef(S.SH, S.SW),
     dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs max pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.c, D.oh, D.ow, D.kh, D.kw)
-  O[D.n, D.c, D.oh, D.ow] = ReduceFn.max(D.kh, D.kw)(
+  O[D.n, D.c, D.oh, D.ow] = ReduceFn.max[D.kh, D.kw](
       TypeFn.cast(
           U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH,
                D.ow * S.SW + D.kw * S.DW,]))
 
 
 @linalg_structured_op
 def pooling_nhwc_min(
     I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
                 S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
     strides=IndexAttrDef(S.SH, S.SW),
     dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs min pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
-  O[D.n, D.oh, D.ow, D.c] = ReduceFn.min(D.kh, D.kw)(
+  O[D.n, D.oh, D.ow, D.c] = ReduceFn.min[D.kh, D.kw](
       TypeFn.cast(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_nhwc_min_unsigned(
     I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
                 S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
     strides=IndexAttrDef(S.SH, S.SW),
     dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs unsigned min pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
-  O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_unsigned(D.kh, D.kw)(
+  O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_unsigned[D.kh, D.kw](
       TypeFn.cast_unsigned(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_ndhwc_sum(
     I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH,
                 S.OW * S.SW + S.KW * S.DW, S.C),
def pooling_ndhwc_max(
     dilations=IndexAttrDef(S.DD, S.DH, S.DW)):
   """Performs 3D max pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
-  O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.max(D.kd, D.kh, D.kw)(
+  O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.max[D.kd, D.kh, D.kw](
       TypeFn.cast(
           U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
                D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_ndhwc_min(
     I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH,
                 S.OW * S.SW + S.KW * S.DW, S.C),
     K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, output=True),
     strides=IndexAttrDef(S.SD, S.SH, S.SW),
     dilations=IndexAttrDef(S.DD, S.DH, S.DW)):
   """Performs 3D min pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
-  O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.min(D.kd, D.kh, D.kw)(
+  O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.min[D.kd, D.kh, D.kw](
       TypeFn.cast(
           U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
                D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def fill_rng_2d(
     min=ScalarDef(F64),

mlir/test/python/dialects/linalg/opdsl/emit_pooling.py

 @linalg_structured_op
 def pooling_max_poly(
     I=TensorDef(T1, S.N, S.H, S.W, S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
     strides=IndexAttrDef(S.SH, S.SW),
     dilations=IndexAttrDef(S.DH, S.DW)):
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
-  O[D.n, D.oh, D.ow, D.c] = ReduceFn.max(D.kh, D.kw)(
+  O[D.n, D.oh, D.ow, D.c] = ReduceFn.max[D.kh, D.kw](
       TypeFn.cast(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_max_unsigned_poly(
     I=TensorDef(T1, S.N, S.H, S.W, S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
     strides=IndexAttrDef(S.SH, S.SW),
     dilations=IndexAttrDef(S.DH, S.DW)):
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
-  O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_unsigned(D.kh, D.kw)(
+  O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_unsigned[D.kh, D.kw](
       TypeFn.cast_unsigned(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_min_poly(
     I=TensorDef(T1, S.N, S.H, S.W, S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
     strides=IndexAttrDef(S.SH, S.SW),
     dilations=IndexAttrDef(S.DH, S.DW)):
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
-  O[D.n, D.oh, D.ow, D.c] = ReduceFn.min(D.kh, D.kw)(
+  O[D.n, D.oh, D.ow, D.c] = ReduceFn.min[D.kh, D.kw](
       TypeFn.cast(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_min_unsigned_poly(
     I=TensorDef(T1, S.N, S.H, S.W, S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
     strides=IndexAttrDef(S.SH, S.SW),
     dilations=IndexAttrDef(S.DH, S.DW)):
   domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
-  O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_unsigned(D.kh, D.kw)(
+  O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_unsigned[D.kh, D.kw](
       TypeFn.cast_unsigned(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 with Context() as ctx, Location.unknown():
   module = Module.create()
   f32 = F32Type.get()
   i32 = IntegerType.get_signless(32)