Files
llvm/mlir/test/python/dialects/transform_structured_ext.py
Ingo Müller 69bc1cbbff [mlir][linalg][transform] Rename {masked_vectorize => vectorize => vectorize_children_and...}. (#66575)
This PR renames the vectorization transform ops as follows:

* `structured.masked_vectorize` => `structured.vectorize`. This reflects
the fact that since [recently](https://reviews.llvm.org/D157774) the op
can also handle the unmasked case.
* `structured.vectorize` =>
`structured.vectorize_children_and_applies_patterns`. This reflects the
fact that the op does not just vectorize the given payload op but all
vectorizable children contained in it, and applies patterns before and
after for preparation and clean-up.

This rename was discussed first
[here](https://reviews.llvm.org/D157774).

The PR also adapts and cleans ups the tablegen description of the
`VectorizeChildrenAndApplyPatternsOp` (formerly `VectorizeOp`).
2023-09-21 15:38:29 +02:00

574 lines
18 KiB
Python

# RUN: %PYTHON %s | FileCheck %s
import functools
from typing import Callable
from mlir.ir import *
from mlir.dialects import transform
from mlir.dialects import pdl
from mlir.dialects.transform import structured
from mlir.dialects.transform import pdl as transform_pdl
def run(f):
with Context(), Location.unknown():
module = Module.create()
with InsertionPoint(module.body):
print("\nTEST:", f.__name__)
f()
module.operation.verify()
print(module)
return f
def create_sequence(func: Callable) -> Callable:
@functools.wraps(func)
def decorated() -> None:
sequence = transform.SequenceOp(
transform.FailurePropagationMode.Propagate,
[],
transform.AnyOpType.get(),
)
with InsertionPoint(sequence.body):
func(sequence.bodyTarget)
transform.YieldOp()
return decorated
@run
@create_sequence
def testBufferizeToAllocationOpCompact(target):
structured.BufferizeToAllocationOp(target)
# CHECK-LABEL: TEST: testBufferizeToAllocationOpCompact
# CHECK: transform.sequence
# CHECK: transform.structured.bufferize_to_allocation
@run
@create_sequence
def testBufferizeToAllocationOpArgs(target):
structured.BufferizeToAllocationOp(
target,
memory_space=3,
memcpy_op="memref.copy",
alloc_op="memref.alloca",
bufferize_destination_only=True,
)
# CHECK-LABEL: TEST: testBufferizeToAllocationOpArgs
# CHECK: transform.sequence
# CHECK: transform.structured.bufferize_to_allocation
# CHECK-SAME: alloc_op = "memref.alloca"
# CHECK-SAME: bufferize_destination_only
# CHECK-SAME: memcpy_op = "memref.copy"
# CHECK-SAME: memory_space = 3
@run
@create_sequence
def testDecompose(target):
structured.DecomposeOp(target)
# CHECK-LABEL: TEST: testDecompose
# CHECK: transform.sequence
# CHECK: transform.structured.decompose
@run
@create_sequence
def testFuseIntoContainingOpTypes(target):
fused = structured.MatchOp.match_op_names(target, ["test.dummy"])
containing = structured.MatchOp.match_op_names(target, ["test.dummy"])
structured.FuseIntoContainingOp(
transform.OperationType.get("test.dummy"),
transform.OperationType.get("test.dummy"),
fused,
containing,
)
# CHECK-LABEL: TEST: testFuseIntoContainingOpTypes
# CHECK: = transform.structured.fuse_into_containing_op
# CHECK-SAME: (!transform.any_op, !transform.any_op) -> (!transform.op<"test.dummy">, !transform.op<"test.dummy">)
@run
@create_sequence
def testFuseIntoContainingOpCompact(target):
fused = structured.MatchOp.match_op_names(target, ["test.dummy"])
containing = structured.MatchOp.match_op_names(target, ["test.dummy"])
structured.FuseIntoContainingOp(fused, containing)
# CHECK-LABEL: TEST: testFuseIntoContainingOpCompact
# CHECK: = transform.structured.fuse_into_containing_op
# CHECK-SAME: (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
@run
@create_sequence
def testGeneralize(target):
structured.GeneralizeOp(target)
# CHECK-LABEL: TEST: testGeneralize
# CHECK: transform.sequence
# CHECK: transform.structured.generalize
@run
@create_sequence
def testInterchange(target):
structured.InterchangeOp(target, iterator_interchange=[1, 0])
# CHECK-LABEL: TEST: testInterchange
# CHECK: transform.sequence
# CHECK: transform.structured.interchange
# CHECK: iterator_interchange = [1, 0]
@run
@create_sequence
def testMapCopyToThreadsOpCompact(target):
structured.MapCopyToThreadsOp(
target, total_num_threads=32, desired_bit_alignment=128
)
# CHECK-LABEL: TEST: testMapCopyToThreadsOpCompact
# CHECK: = transform.structured.gpu.map_copy_to_threads
# CHECK-SAME: total_num_threads = 32
# CHECK-SAME: desired_bit_alignment = 128
# CHECK-SAME: (!transform.any_op) -> (!transform.any_op, !transform.any_op)
@run
@create_sequence
def testMapCopyToThreadsOpTypes(target):
structured.MapCopyToThreadsOp(
transform.OperationType.get("test.opA"),
transform.OperationType.get("test.opB"),
target,
total_num_threads=32,
desired_bit_alignment=128,
)
# CHECK-LABEL: TEST: testMapCopyToThreadsOpTypes
# CHECK: = transform.structured.gpu.map_copy_to_threads
# CHECK-SAME: total_num_threads = 32
# CHECK-SAME: desired_bit_alignment = 128
# CHECK-SAME: (!transform.any_op) -> (!transform.op<"test.opA">, !transform.op<"test.opB">)
@run
@create_sequence
def testMatchOpNamesString(target):
structured.MatchOp.match_op_names(target, "test.dummy")
# CHECK-LABEL: TEST: testMatchOpNamesString
# CHECK: transform.structured.match ops
# CHECK-SAME: ["test.dummy"]
# CHECK-SAME: (!transform.any_op) -> !transform.any_op
@run
@create_sequence
def testMatchOpNamesList(target):
structured.MatchOp.match_op_names(target, ["test.dummy"])
# CHECK-LABEL: TEST: testMatchOpNamesList
# CHECK: transform.structured.match ops
# CHECK-SAME: ["test.dummy"]
# CHECK-SAME: (!transform.any_op) -> !transform.any_op
@run
@create_sequence
def testVectorizeNoArgs(target):
structured.VectorizeOp(target)
# CHECK-LABEL: TEST: testVectorizeNoArgs
# CHECK: transform.sequence
# CHECK: transform.structured.vectorize
# CHECK-NOT: vector_sizes
def testVectorizeStatic(target):
structured.VectorizeOp(target, [16, 4])
# CHECK-LABEL: TEST: testVectorizeStatic
# CHECK: transform.sequence
# CHECK: transform.structured.vectorize
# CHECK-SAME: vector_sizes [16, 4]
@run
@create_sequence
def testVectorizeArray(target):
sizes = Attribute.parse("[16, 4]")
structured.VectorizeOp(target, sizes)
# CHECK-LABEL: TEST: testVectorizeArray
# CHECK: transform.sequence
# CHECK: transform.structured.vectorize
# CHECK-SAME: vector_sizes [16, 4]
@run
@create_sequence
def testVectorizeMixed(target):
sz1 = structured.MatchOp.match_op_names(target, ["arith.constant"])
sz2 = Attribute.parse("4")
structured.VectorizeOp(target, [sz1, sz2])
# CHECK-LABEL: TEST: testVectorizeMixed
# CHECK: transform.sequence
# CHECK: %[[V0:.*]] = transform.structured.match
# CHECK: transform.structured.vectorize
# CHECK-SAME: vector_sizes [%[[V0]] : !transform.any_op, 4]
@run
@create_sequence
def testVectorizeScalable(target):
sz1 = structured.MatchOp.match_op_names(target, ["arith.constant"])
sz2 = Attribute.parse("4")
structured.VectorizeOp(target, [16, [sz1], [sz2], [8]])
# CHECK-LABEL: TEST: testVectorizeScalable
# CHECK: transform.sequence
# CHECK-DAG: %[[V0:.*]] = transform.structured.match
# CHECK-DAG: transform.structured.vectorize
# CHECK-SAME: vector_sizes [16, [%[[V0]] : !transform.any_op], [4], [8]]
@run
@create_sequence
def testVectorizeArgs(target):
structured.VectorizeOp(target, [16, 4], vectorize_nd_extract=True)
# CHECK-LABEL: TEST: testVectorizeArgs
# CHECK: transform.sequence
# CHECK: transform.structured.vectorize
# CHECK-SAME: vectorize_nd_extract
@run
@create_sequence
def testMatchOpNamesTyped(target):
structured.MatchOp.match_op_names(
transform.OperationType.get("test.dummy"),
target,
["test.dummy"],
)
# CHECK-LABEL: TEST: testMatchOpNamesTyped
# CHECK: transform.structured.match ops
# CHECK-SAME: ["test.dummy"]
# CHECK-SAME: (!transform.any_op) -> !transform.op<"test.dummy">
@run
@create_sequence
def testMultitileSizesCompact(target):
structured.MultiTileSizesOp(
transform.AnyOpType.get(), target, dimension=1, target_size=42
)
# CHECK-LABEL: TEST: testMultitileSizes
# CHECK: transform.sequence
# CHECK-NOT: divisor
# CHECK: transform.structured.multitile_sizes
# CHECK-NOT: divisor
# CHECK-DAG: dimension = 1
# CHECK-NOT: divisor
# CHECK-DAG: target_size = 42
# CHECK-NOT: divisor
@run
@create_sequence
def testMultitileSizesAllArgs(target):
structured.MultiTileSizesOp(
transform.AnyOpType.get(),
target,
dimension=1,
target_size=42,
divisor=2,
)
# CHECK-LABEL: TEST: testMultitileSizes
# CHECK: transform.sequence
# CHECK: transform.structured.multitile_sizes
# CHECK-DAG: dimension = 1
# CHECK-DAG: divisor = 2
# CHECK-DAG: target_size = 42
@run
@create_sequence
def testPadOpNoArgs(target):
structured.PadOp(target)
# CHECK-LABEL: TEST: testPadOpNoArgs
# CHECK: transform.sequence
# CHECK: transform.structured.pad
# CHECK-NOT: copy_back_op
# CHECK-NOT: pack_paddings
# CHECK-NOT: pad_to_multiple_of
# CHECK-NOT: padding_dimensions
# CHECK-NOT: padding_values
# CHECK-NOT: transpose_paddings
@run
@create_sequence
def testPadOpArgs(target):
structured.PadOp(
target,
padding_values=[FloatAttr.get_f32(42.0), StringAttr.get("0")],
padding_dimensions=Attribute.parse("[1]"),
pad_to_multiple_of=[128],
pack_paddings=[0],
transpose_paddings=[[1, Attribute.parse("0")], Attribute.parse("[0, 1]")],
copy_back_op="linalg.copy",
)
# CHECK-LABEL: TEST: testPadOpArgs
# CHECK: transform.sequence
# CHECK: transform.structured.pad
# CHECK-DAG: copy_back_op = "linalg.copy"
# CHECK-DAG: pack_paddings = [0]
# CHECK-DAG: pad_to_multiple_of = [128]
# CHECK-DAG: padding_dimensions = [1]
# CHECK-DAG: padding_values = [4.200000e+01 : f32, "0"]
# CHECK-DAG: transpose_paddings = {{\[}}[1, 0], [0, 1]]
@run
@create_sequence
def testScalarize(target):
structured.ScalarizeOp(target)
# CHECK-LABEL: TEST: testScalarize
# CHECK: transform.structured.scalarize
@run
@create_sequence
def testSplit(target):
split = structured.SplitOp(target, dimension=1, split_point=42)
structured.SplitOp(split.results[0], dimension=3, split_point=split.results[1])
# CHECK-LABEL: TEST: testSplit
# CHECK: %[[F:.+]], %[[S:.+]] = transform.structured.split %{{.*}} after 42 {dimension = 1
# CHECK: transform.structured.split %[[F]] after %[[S]] {dimension = 3
@run
@create_sequence
def testTileCompact(target):
structured.TileOp(target, sizes=[4, 8], interchange=[0, 1])
# CHECK-LABEL: TEST: testTileCompact
# CHECK: transform.sequence
# CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile %{{.*}}[4, 8]
# CHECK: interchange = [0, 1]
@run
@create_sequence
def testTileAttributes(target):
attr = DenseI64ArrayAttr.get([4, 8])
ichange = DenseI64ArrayAttr.get([0, 1])
structured.TileOp(target, sizes=attr, interchange=ichange)
# CHECK-LABEL: TEST: testTileAttributes
# CHECK: transform.sequence
# CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile %{{.*}}[4, 8]
# CHECK: interchange = [0, 1]
@run
@create_sequence
def testTileZero(target):
structured.TileOp(target, sizes=[4, 0, 2, 0], interchange=[0, 1, 2, 3])
# CHECK-LABEL: TEST: testTileZero
# CHECK: transform.sequence
# CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile %{{.*}}[4, 0, 2, 0]
# CHECK: interchange = [0, 1, 2, 3]
@run
def testTileDynamic():
with_pdl = transform_pdl.WithPDLPatternsOp(pdl.OperationType.get())
with InsertionPoint(with_pdl.body):
sequence = transform.SequenceOp(
transform.FailurePropagationMode.Propagate, [], with_pdl.bodyTarget
)
with InsertionPoint(sequence.body):
m1 = transform_pdl.PDLMatchOp(
pdl.OperationType.get(), sequence.bodyTarget, "first"
)
m2 = transform_pdl.PDLMatchOp(
pdl.OperationType.get(), sequence.bodyTarget, "second"
)
structured.TileOp(sequence.bodyTarget, sizes=[m1, 3, m2, 0])
transform.YieldOp()
# CHECK-LABEL: TEST: testTileDynamic
# CHECK: %[[FIRST:.+]] = pdl_match
# CHECK: %[[SECOND:.+]] = pdl_match
# CHECK: %{{.+}}, %{{.+}}:3 = transform.structured.tile %{{.*}}[%[[FIRST]], 3, %[[SECOND]], 0]
@run
@create_sequence
def testTileExplicitLoopTypeSingle(target):
structured.TileOp(transform.OperationType.get("scf.for"), target, sizes=[2, 3, 4])
# CHECK-LABEL: TEST: testTileExplicitLoopTypeSingle
# CHECK: = transform.structured.tile %{{.*}} : (!{{.*}}) ->
# CHECK-COUNT-3: !transform.op<"scf.for">
@run
@create_sequence
def testTileExplicitLoopTypeAll(target):
types = [
transform.OperationType.get(x)
for x in ["scf.for", "scf.parallel", "scf.forall"]
]
structured.TileOp(types, target, sizes=[2, 3, 4])
# 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
@create_sequence
def testTileScalable(target):
structured.TileOp(
target,
sizes=[4, [2]],
)
# CHECK-LABEL: TEST: testTileScalable
# CHECK: transform.sequence
# CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile %{{.*}}[4, [2]]
@run
@create_sequence
def testTileToForallCompact(target):
matmul = transform.CastOp(transform.OperationType.get("linalg.matmul"), target)
structured.TileToForallOp(matmul, num_threads=[2, 3, 4])
# 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.any_op, !transform.any_op)
@run
@create_sequence
def testTileToForallLoopsAndTileOpTypes(target):
structured.TileToForallOp(
transform.OperationType.get("scf.forall"), # loops_type
transform.OperationType.get("linalg.matmul"), # tiled_op_type
target,
num_threads=[2, 3, 4],
)
# 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.op<"scf.forall">, !transform.op<"linalg.matmul">)
@run
@create_sequence
def testTileToForallTileSizes(target):
structured.TileToForallOp(target, tile_sizes=[2, 3, 4])
# CHECK-LABEL: TEST: testTileToForallTileSizes
# CHECK: = transform.structured.tile_to_forall_op
# CHECK-SAME: num_threads [] tile_sizes [2, 3, 4]
@run
@create_sequence
def testTileToForallMixedDynamic(target):
n = structured.MatchOp.match_op_names(target, ["test.dummy"])
structured.TileToForallOp(target, num_threads=[n, 3, 4])
# CHECK-LABEL: TEST: testTileToForallMixedDynamic
# CHECK: = transform.structured.tile_to_forall_op
# CHECK-SAME: num_threads [%{{.*}} : !transform.any_op, 3, 4]
@run
@create_sequence
def testTileToForallPackedDynamic(target):
n = structured.MatchOp.match_op_names(target, ["test.dummy"])
structured.TileToForallOp(target, num_threads=n)
# CHECK-LABEL: TEST: testTileToForallPackedDynamic
# CHECK: = transform.structured.tile_to_forall_op
# CHECK-SAME: num_threads *(%0 : !transform.any_op)
@run
@create_sequence
def testTileToForallMapping(target):
mapping = Attribute.parse("[ #gpu.thread<y>, #gpu.thread<x> ]")
structured.TileToForallOp(target, num_threads=[2, 3], mapping=mapping)
# CHECK-LABEL: TEST: testTileToForallMapping
# CHECK: = transform.structured.tile_to_forall_op
# CHECK-SAME: mapping = [#gpu.thread<y>, #gpu.thread<x>]
@run
@create_sequence
def testVectorizeChildrenAndApplyPatternsAllAttrs(target):
structured.VectorizeChildrenAndApplyPatternsOp(
target,
disable_multi_reduction_to_contract_patterns=True,
disable_transfer_permutation_map_lowering_patterns=True,
vectorize_nd_extract=True,
vectorize_padding=True,
)
# CHECK-LABEL: TEST: testVectorizeChildrenAndApplyPatternsAllAttrs
# CHECK: transform.sequence
# CHECK: = transform.structured.vectorize
# CHECK-SAME: disable_multi_reduction_to_contract_patterns
# CHECK-SAME: disable_transfer_permutation_map_lowering_patterns
# CHECK-SAME: vectorize_nd_extract
# CHECK-SAME: vectorize_padding
@run
@create_sequence
def testVectorizeChildrenAndApplyPatternsNoAttrs(target):
structured.VectorizeChildrenAndApplyPatternsOp(
target,
disable_multi_reduction_to_contract_patterns=False,
disable_transfer_permutation_map_lowering_patterns=False,
vectorize_nd_extract=False,
vectorize_padding=False,
)
# CHECK-LABEL: TEST: testVectorizeChildrenAndApplyPatternsNoAttrs
# CHECK: transform.sequence
# CHECK: = transform.structured.vectorize
# CHECK-NOT: disable_multi_reduction_to_contract_patterns
# CHECK-NOT: disable_transfer_permutation_map_lowering_patterns
# CHECK-NOT: vectorize_nd_extract
# CHECK-NOT: vectorize_padding
@run
@create_sequence
def testMatchInterfaceEnum(target):
names = ArrayAttr.get([StringAttr.get("test.dummy")])
result_type = transform.AnyOpType.get()
fused = structured.MatchOp.__base__(
result_type,
target,
ops=names,
interface=structured.MatchInterfaceEnum.LinalgOp,
)
# CHECK-LABEL: TEST: testMatchInterfaceEnum
# CHECK: transform.sequence
# CHECK: = transform.structured.match
# CHECK: interface{LinalgOp}
@run
@create_sequence
def testMatchInterfaceEnumReplaceAttributeBuilder(target):
@register_attribute_builder("MatchInterfaceEnum", replace=True)
def match_interface_enum(x, context):
if x == "LinalgOp":
y = 0
elif x == "TilingInterface":
y = 1
return IntegerAttr.get(IntegerType.get_signless(32, context=context), y)
names = ArrayAttr.get([StringAttr.get("test.dummy")])
result_type = transform.AnyOpType.get()
fused = structured.MatchOp.__base__(
result_type,
target,
ops=names,
interface="TilingInterface",
)
# CHECK-LABEL: TEST: testMatchInterfaceEnumReplaceAttributeBuilder
# CHECK: transform.sequence
# CHECK: = transform.structured.match
# CHECK: interface{TilingInterface}