[mlir][transform][linalg][python] Add extended TileToForallOp.

[ingomueller-net]

Summary

This patch adds a mixin for TileToForallOp to
_structured_transform_ops_ext.py with syntactic sugar for construction
such ops. First, the types of the results are made optional and filled
with common default values if omitted. Second, for num_threads and
tile_sizes, the three possible forms (static, dynamic, or packed), can
now all be given through the same respective argument, which gets
dispatched to the correct form-specific argument automatically.

[ingomueller-net]

Please verify the following assumptions that went into the design:

1. Specifying num_threads is the common case and tile_sizes is not. This would justify the current design: former is a positional argument whereas the latter is a keyword-only argument. (Note that they can't be positional both at the same time because they have exactly the same possible types.)
2. There are actually cases where we want to change the types of forall_op and tiled_op and there are several respective choices. If there no reasonable cases, then we can get away with fewer overloads and use the current default types always. If there are just two cases (the current default and "any type"), then a Boolean flag might be more appropriate.
3. It is more common to change the type of forall_op than it is to change the type of tiled_op. If it is the other way around, it may be better to flip their order in the overloads.

[ftynse]

mlir/python/mlir/dialects/_structured_transform_ops_ext.py

430–432

I suppose this may be useful to other ops. Could we make it a standalone function from the start?

438–441

Nit: it feels a bit magical that an operation is interpreted as packed and a singleton list of operations is interpreted as one value, but we can try and see how often people are confused with this. (In C++, I'd wrap them into different "tag" structures, but Python is more tolerant towards such type tricks).

[ftynse]

Specifying num_threads is the common case and tile_sizes is not. This would justify the current design: former is a positional argument whereas the latter is a keyword-only argument. (Note that they can't be positional both at the same time because they have exactly the same possible types.)

Per https://peps.python.org/pep-0020/ "explicit is better than implicit", so I would strongly suggest to have both as keyword arguments.

There are actually cases where we want to change the types of forall_op and tiled_op and there are several respective choices. If there no reasonable cases, then we can get away with fewer overloads and use the current default types always. If there are just two cases (the current default and "any type"), then a Boolean flag might be more appropriate.

There aren't that many cases that use the types extensively, but I'd argue that there should be. The point of the type system is to specify visible static preconditions on the ops, e.g., tiling only applies to tilable ops (though we don't have a type for this). The result type of a transform op is by default very open so the users can supply whatever type the following op needs and avoid a cast. I'd go with any_op as a default in case of the missing result type.

It is more common to change the type of forall_op than it is to change the type of tiled_op. If it is the other way around, it may be better to flip their order in the overloads.

The order *must* follow the order of results in the operation. Anything else without explicit keywords would be extremely error-prone. Furthermore, tile_to_forall already uses a different order of loop/op results than the "regular" tile. That should be fixed, but it shouldn't be your problem (fixes are certainly welcome).

[ftynse]

The order *must* follow the order of results in the operation. Anything else without explicit keywords would be extremely error-prone. Furthermore, tile_to_forall already uses a different order of loop/op results than the "regular" tile. That should be fixed, but it shouldn't be your problem (fixes are certainly welcome).

The order currently doesn't match, let's use None for one of the arguments if need be, but not flip the order. Requesting changes for this one.

[ingomueller-net]

In D155090#4505128, @ftynse wrote:

The order *must* follow the order of results in the operation. Anything else without explicit keywords would be extremely error-prone. Furthermore, tile_to_forall already uses a different order of loop/op results than the "regular" tile. That should be fixed, but it shouldn't be your problem (fixes are certainly welcome).

The order currently doesn't match, let's use None for one of the arguments if need be, but not flip the order. Requesting changes for this one.

Yeah, makes sense, will do. What about the overload with just one type, though? I am thinking that maybe it shouldn't exist: either you specify all types (in __init__ order) or none. Otherwise, what's the most intuitive/import of the two to specify?

[ftynse]

In D155090#4509347, @ingomueller-net wrote:

In D155090#4505128, @ftynse wrote:

The order *must* follow the order of results in the operation. Anything else without explicit keywords would be extremely error-prone. Furthermore, tile_to_forall already uses a different order of loop/op results than the "regular" tile. That should be fixed, but it shouldn't be your problem (fixes are certainly welcome).

The order currently doesn't match, let's use None for one of the arguments if need be, but not flip the order. Requesting changes for this one.

Yeah, makes sense, will do. What about the overload with just one type, though? I am thinking that maybe it shouldn't exist: either you specify all types (in __init__ order) or none. Otherwise, what's the most intuitive/import of the two to specify?

I'd always specify both for now. People either care about types or don't use them at all.

[ingomueller-net]

Addressing comments from @ftynse:

• Make _dispatch_mixed_values a free function.
• Make num_threads a kewword-only argument.
• Re-order tiled_op_type and loops_type to match the constructor/op definition.
• Make the default types AnyOpTypes.

I also did the following changes:

• Remove the overload with only one type. I don't see how a user could predict which of the two is meant if only one of them is given, so this seems error-prone. Also, if need be, such an overload could be added later.
• Format the changed lines with yapf.
• Remove a spurious print.

[ftynse]

Format the changed lines with yapf.

We are now supposed to format with black: https://discourse.llvm.org/t/rfc-document-and-standardize-python-code-style/68257.

[ingomueller-net]

More changes:

• Reformat changed lines with black.
• Change commit title.

[ingomueller-net]

mlir/python/mlir/dialects/_structured_transform_ops_ext.py

438–441

I am not a Python guru but I thought that this is a usual pattern. For example, the built-in open function accepts either a "path-like object", which in turn is an str, a bytes, or an os.Path, or an integer corresponding to a Unix file descriptor. All forms somehow identify the file that should be opened and the implementation of the function (supposedly) does some dispatching to interpret it.

mlir/python/mlir/dialects/_structured_transform_ops_ext.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
 
 try:
     from ..ir import *
     from ._ods_common import get_op_result_or_value as _get_op_result_or_value
     from ..dialects import pdl, transform
 except ImportError as e:
     raise RuntimeError("Error loading imports from extension module") from e
 
-from typing import List, Optional, Sequence, Union, overload
+from typing import List, Optional, Sequence, Tuple, Union, overload
 
 IntOrAttrList = Sequence[Union[IntegerAttr, int]]
 OptionalIntList = Optional[Union[ArrayAttr, IntOrAttrList]]
 
 BoolOrAttrList = Sequence[Union[BoolAttr, bool]]
 OptionalBoolList = Optional[Union[ArrayAttr, BoolOrAttrList]]
 
 
super().__init__(
         )
 
   def __extract_values(self, attr: Optional[DenseI64ArrayAttr]) -> List[int]:
     if not attr:
       return []
     return [element for element in attr]
 
 
+class TileToForallOp:
+  """Specialization for TileToForallOp class."""
+  _MixedValues = Union[Sequence[Union[int, IntegerAttr, Operation, Value]],
+                       ArrayAttr, Operation, Value, OpView]
+
+  @overload
+  def __init__(
+        self,
+        tiled_op_type: Type,
+        loops_type: Type,
+        target: Union[Operation, Value, OpView],
+        num_threads: Optional[_MixedValues] = None,
+        *,
+        tile_sizes: _MixedValues = None,
+        mapping=None,
+        loc=None,
+        ip=None,
+    ):
+    ...
+
+  @overload
+  def __init__(
+        self,
+        loops_type: Type,
+        target: Union[Operation, Value, OpView],
+        num_threads: Optional[_MixedValues] = None,
+        *,
+        tile_sizes: _MixedValues = None,
+        mapping=None,
+        loc=None,
+        ip=None,
+    ):
+    ...
+
+  @overload
+  def __init__(
+        self,
+        target: Union[Operation, Value, OpView],
+        num_threads: Optional[_MixedValues] = None,
+        *,
+        tile_sizes: _MixedValues = None,
+        mapping=None,
+        loc=None,
+        ip=None,
+    ):
+    ...
+
+  def __init__(
+        self,
+        arg0: Union[Type,                              # tile_op_type/loops_type
+                    Union[Operation, Value, OpView]],  # target
+        arg1: Union[Type,                              # loops_type
+                    Union[Operation, Value, OpView],   # target
+                    Optional[_MixedValues]] = None,    # num_threads
+        arg2: Optional[Union[
+                Union[Operation, Value, OpView],       # target
+                      Optional[_MixedValues]]] = None, # num_threads
+        arg3: Optional[_MixedValues] = None,           # num_threads
+        *,
+        tile_sizes: _MixedValues = None,
+        mapping=None,
+        loc=None,
+        ip=None,
+    ):
+
+    # `Type` arguments in the front are optional: add default values to front.
+    if isinstance(arg1, Type):
+      # First overload: all arguments provided.
+      args = [arg0, arg1, arg2, arg3]
+    elif isinstance(arg0, Type):
+      # Second overload: tiled_op_type missing.
+      target = arg1
+      tiled_op_type = target.type
+      args = [tiled_op_type, arg0, arg1, arg2]
+    else:
+      # Last overload: both types missing.
+      loops_type = transform.OperationType.get('scf.forall')
+      target = arg0
+      tiled_op_type = target.type
+      args = [tiled_op_type, loops_type, arg0, arg1]
+    (tiled_op_type, loops_type, target, num_threads) = args
+    print(args)
+
+    # Unpack mixed num_threads.
+    (dynamic_num_threads,
+     packed_num_threads,
+     num_threads_attr) = self._dispatch_mixed_values(num_threads)
+
+    # Unpack mixed tile_sizes.
+    (dynamic_tile_sizes,
+     packed_tile_sizes,
+     tile_sizes_attr) = self._dispatch_mixed_values(tile_sizes)
+
+    super().__init__(
+            loops_type,
+            tiled_op_type,
+            target=target,
+            tile_sizes=dynamic_tile_sizes,
+            packed_tile_sizes=packed_tile_sizes,
+            static_tile_sizes=tile_sizes_attr,
+            num_threads=dynamic_num_threads,
+            packed_num_threads=packed_num_threads,
+            static_num_threads=num_threads_attr,
+            mapping=mapping,
+            loc=loc,
+            ip=ip,
+        )
+
+  @classmethod
+  def _dispatch_mixed_values(cls: Type, values: _MixedValues) \
+      -> Tuple[List[Value], Union[Operation, Value, OpView], DenseI64ArrayAttr]:
+    dynamic_values = []
+    packed_values = None
+    attr_value = None
+    if isinstance(values, ArrayAttr):
+      attr_value = values
+    elif isinstance(values, (Operation, Value, OpView)):
+      packed_values = values
+    else:
+      static_values = []
+      for size in values or []:
+        if isinstance(size, int):
+          static_values.append(size)
+        else:
+          static_values.append(ShapedType.get_dynamic_size())
+          dynamic_values.append(_get_op_result_or_value(size))
+      attr_value = DenseI64ArrayAttr.get(static_values)
+
+    return (dynamic_values, packed_values, attr_value)
+
+
 class VectorizeOp:
     """Specialization for VectorizeOp class."""
 
     def __init__(
         self,
         target: Union[Operation, Value],
         *,
         vectorize_padding: Union[bool, BoolAttr] = False,

mlir/test/python/dialects/transform_structured_ext.py

with InsertionPoint(sequence.body):
         transform.YieldOp()
     # CHECK-LABEL: TEST: testTileExplicitLoopTypeAll
     # CHECK: = transform.structured.tile
     # CHECK-SAME : (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">,
     # CHECK-SAME: !transform.op<"scf.parallel">, !transform.op<"scf.forall">
 
 
 @run
+def testTileToForallCompact():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.PROPAGATE, [],
+        transform.OperationType.get('linalg.matmul')
+    )
+    with InsertionPoint(sequence.body):
+        structured.TileToForallOp(sequence.bodyTarget, [2, 3, 4])
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testTileToForallCompact
+    # CHECK: = transform.structured.tile_to_forall_op
+    # CHECK-SAME: num_threads [2, 3, 4] tile_sizes []
+    # CHECK-SAME: (!transform.op<"linalg.matmul">) -> (!transform.op<"scf.forall">, !transform.op<"linalg.matmul">)
+
+
+@run
+def testTileToForallLoopsType():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get()
+    )
+    with InsertionPoint(sequence.body):
+        structured.TileToForallOp(
+            transform.AnyOpType.get(), # loops_type
+            sequence.bodyTarget, [2, 3, 4],
+        )
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testTileToForallLoopsType
+    # CHECK: = transform.structured.tile_to_forall_op
+    # CHECK-SAME: num_threads [2, 3, 4] tile_sizes []
+    # CHECK-SAME: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+
+
+@run
+def testTileToForallLoopsAndTileOpTypes():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get()
+    )
+    with InsertionPoint(sequence.body):
+        structured.TileToForallOp(
+            transform.OperationType.get('linalg.matmul'),  # tiled_op_type
+            transform.AnyOpType.get(),                     # loops_type
+            sequence.bodyTarget, [2, 3, 4],
+        )
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testTileToForallLoopsAndTileOpTypes
+    # CHECK: = transform.structured.tile_to_forall_op
+    # CHECK-SAME: num_threads [2, 3, 4] tile_sizes []
+    # CHECK-SAME: (!transform.any_op) -> (!transform.any_op, !transform.op<"linalg.matmul">)
+
+
+@run
+def testTileToForallTileSizes():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get()
+    )
+    with InsertionPoint(sequence.body):
+        structured.TileToForallOp(sequence.bodyTarget, tile_sizes=[2, 3, 4])
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testTileToForallTileSizes
+    # CHECK: = transform.structured.tile_to_forall_op
+    # CHECK-SAME: num_threads [] tile_sizes [2, 3, 4]
+
+
+@run
+def testTileToForallMixedDynamic():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get()
+    )
+    with InsertionPoint(sequence.body):
+        sz = structured.MatchOp.match_op_names(sequence.bodyTarget, ['test.dummy'])
+        structured.TileToForallOp(sequence.bodyTarget, [sz, 3, 4])
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testTileToForallMixedDynamic
+    # CHECK: = transform.structured.tile_to_forall_op
+    # CHECK-SAME: num_threads [%{{.*}} : !pdl.operation, 3, 4]
+
+
+@run
+def testTileToForallMPackedDynamic():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get()
+    )
+    with InsertionPoint(sequence.body):
+        sz = structured.MatchOp.match_op_names(sequence.bodyTarget, ['test.dummy'])
+        structured.TileToForallOp(sequence.bodyTarget, sz)
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testTileToForallMPackedDynamic
+    # CHECK: = transform.structured.tile_to_forall_op
+    # CHECK-SAME: num_threads *(%0 : !pdl.operation)
+
+
+@run
+def testTileToForallMapping():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get()
+    )
+    with InsertionPoint(sequence.body):
+        mapping = Attribute.parse('[ #gpu.thread<y>, #gpu.thread<x> ]')
+        structured.TileToForallOp(
+            sequence.bodyTarget, [2, 3],
+            mapping=mapping
+        )
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testTileToForallMapping
+    # CHECK: = transform.structured.tile_to_forall_op
+    # CHECK-SAME: mapping = [#gpu.thread<y>, #gpu.thread<x>]
+
+
+@run
 def testVectorize():
     sequence = transform.SequenceOp(
         transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get()
     )
     with InsertionPoint(sequence.body):
         structured.VectorizeOp(sequence.bodyTarget, vectorize_padding=True)
         transform.YieldOp()
     # CHECK-LABEL: TEST: testVectorize
     # CHECK: transform.sequence
     # CHECK: = transform.structured.vectorize
     # CHECK: {vectorize_padding}