2c1c67674c
This patch is a first pass at making consistent syntax across the `LinalgTransformOp`s that use dynamic index lists for size parameters. Previously, there were two different forms: inline types in the list, or place them in the functional style tuple. This patch goes for the latter. In order to do this, the `printPackedOrDynamicIndexList`, `printDynamicIndexList` and their `parse` counterparts were modified so that the types can be optionally provided to the corresponding custom directives. All affected ops now use tablegen `assemblyFormat`, so custom `parse`/`print` functions have been removed. There are a couple ops that will likely add dynamic size support, and once that happens it should be made sure that the assembly remains consistent with the changes in this patch. The affected ops are as follows: `pack`, `pack_greedily`, `tile_using_forall`. The `tile_using_for` and `vectorize` ops already used this syntax, but their custom assembly was removed. --------- Co-authored-by: Oleksandr "Alex" Zinenko <ftynse@gmail.com>
605 lines
19 KiB
Python
605 lines
19 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
|
|
from mlir.dialects.transform.extras import constant_param
|
|
|
|
|
|
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
|
|
|
|
|
|
@run
|
|
@create_sequence
|
|
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]], 4]
|
|
|
|
|
|
@run
|
|
@create_sequence
|
|
def testVectorizeEmpty(target):
|
|
structured.VectorizeOp(target, [])
|
|
# CHECK-LABEL: TEST: testVectorizeEmpty
|
|
# CHECK: transform.sequence
|
|
# CHECK: transform.structured.vectorize
|
|
# CHECK-NOT: vector_sizes
|
|
|
|
|
|
@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]]], [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,
|
|
pad_to_multiple_of=[128],
|
|
padding_values=[FloatAttr.get_f32(42.0), StringAttr.get("0")],
|
|
padding_dimensions=Attribute.parse("[1]"),
|
|
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: pad_to_multiple_of [128]
|
|
# CHECK-DAG: copy_back_op = "linalg.copy"
|
|
# CHECK-DAG: pack_paddings = [0]
|
|
# 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 testPadOpArgsParam(target):
|
|
structured.PadOp(
|
|
target,
|
|
pad_to_multiple_of=[constant_param(128), Attribute.parse("2"), 10],
|
|
padding_dimensions=Attribute.parse("[0, 1, 2]"),
|
|
)
|
|
# CHECK-LABEL: TEST: testPadOpArgsParam
|
|
# CHECK: transform.sequence
|
|
# CHECK-DAG: %[[P:.*]] = transform.param.constant 128
|
|
# CHECK: transform.structured.pad
|
|
# CHECK-DAG: pad_to_multiple_of [%[[P]], 2, 10]
|
|
# CHECK-DAG: padding_dimensions = [0, 1, 2]
|
|
|
|
|
|
@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.TileUsingForOp(target, sizes=[4, 8], interchange=[0, 1])
|
|
# CHECK-LABEL: TEST: testTileCompact
|
|
# CHECK: transform.sequence
|
|
# CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile_using_for %{{.*}}[4, 8]
|
|
# CHECK: interchange = [0, 1]
|
|
|
|
|
|
@run
|
|
@create_sequence
|
|
def testTileAttributes(target):
|
|
attr = DenseI64ArrayAttr.get([4, 8])
|
|
ichange = DenseI64ArrayAttr.get([0, 1])
|
|
structured.TileUsingForOp(target, sizes=attr, interchange=ichange)
|
|
# CHECK-LABEL: TEST: testTileAttributes
|
|
# CHECK: transform.sequence
|
|
# CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile_using_for %{{.*}}[4, 8]
|
|
# CHECK: interchange = [0, 1]
|
|
|
|
|
|
@run
|
|
@create_sequence
|
|
def testTileZero(target):
|
|
structured.TileUsingForOp(target, sizes=[4, 0, 2, 0], interchange=[0, 1, 2, 3])
|
|
# CHECK-LABEL: TEST: testTileZero
|
|
# CHECK: transform.sequence
|
|
# CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile_using_for %{{.*}}[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.TileUsingForOp(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_using_for %{{.*}}[%[[FIRST]], 3, %[[SECOND]], 0]
|
|
|
|
|
|
@run
|
|
@create_sequence
|
|
def testTileExplicitLoopTypeSingle(target):
|
|
structured.TileUsingForOp(
|
|
transform.OperationType.get("scf.for"), target, sizes=[2, 3, 4]
|
|
)
|
|
# CHECK-LABEL: TEST: testTileExplicitLoopTypeSingle
|
|
# CHECK: = transform.structured.tile_using_for %{{.*}} : (!{{.*}}) ->
|
|
# 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.TileUsingForOp(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.TileUsingForOp(
|
|
target,
|
|
sizes=[4, [2]],
|
|
)
|
|
# CHECK-LABEL: TEST: testTileScalable
|
|
# CHECK: transform.sequence
|
|
# CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile_using_for %{{.*}}[4, [2]]
|
|
|
|
|
|
@run
|
|
@create_sequence
|
|
def testTileToForallCompact(target):
|
|
matmul = transform.CastOp(transform.OperationType.get("linalg.matmul"), target)
|
|
structured.TileUsingForallOp(matmul, num_threads=[2, 3, 4])
|
|
# CHECK-LABEL: TEST: testTileToForallCompact
|
|
# CHECK: = transform.structured.tile_using_forall
|
|
# CHECK-SAME: num_threads [2, 3, 4]
|
|
# CHECK-SAME: (!transform.op<"linalg.matmul">) -> (!transform.any_op, !transform.any_op)
|
|
|
|
|
|
@run
|
|
@create_sequence
|
|
def testTileToForallLoopsAndTileOpTypes(target):
|
|
structured.TileUsingForallOp(
|
|
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_using_forall
|
|
# CHECK-SAME: num_threads [2, 3, 4]
|
|
# CHECK-SAME: (!transform.any_op) -> (!transform.op<"scf.forall">, !transform.op<"linalg.matmul">)
|
|
|
|
|
|
@run
|
|
@create_sequence
|
|
def testTileToForallTileSizes(target):
|
|
structured.TileUsingForallOp(target, tile_sizes=[2, 3, 4])
|
|
# CHECK-LABEL: TEST: testTileToForallTileSizes
|
|
# CHECK: = transform.structured.tile_using_forall
|
|
# CHECK-SAME: tile_sizes [2, 3, 4]
|
|
|
|
|
|
@run
|
|
@create_sequence
|
|
def testTileToForallMixedDynamic(target):
|
|
n = structured.MatchOp.match_op_names(target, ["test.dummy"])
|
|
structured.TileUsingForallOp(target, num_threads=[n, 3, 4])
|
|
# CHECK-LABEL: TEST: testTileToForallMixedDynamic
|
|
# CHECK: = transform.structured.tile_using_forall
|
|
# CHECK-SAME: num_threads [%{{.*}}, 3, 4] : (!transform.any_op, !transform.any_op)
|
|
|
|
|
|
@run
|
|
@create_sequence
|
|
def testTileToForallPackedDynamic(target):
|
|
n = structured.MatchOp.match_op_names(target, ["test.dummy"])
|
|
structured.TileUsingForallOp(target, num_threads=n)
|
|
# CHECK-LABEL: TEST: testTileToForallPackedDynamic
|
|
# CHECK: = transform.structured.tile_using_forall
|
|
# CHECK-SAME: num_threads *(%0) : (!transform.any_op, !transform.any_op)
|
|
|
|
|
|
@run
|
|
@create_sequence
|
|
def testTileToForallMapping(target):
|
|
mapping = Attribute.parse("[ #gpu.thread<y>, #gpu.thread<x> ]")
|
|
structured.TileUsingForallOp(target, num_threads=[2, 3], mapping=mapping)
|
|
# CHECK-LABEL: TEST: testTileToForallMapping
|
|
# CHECK: = transform.structured.tile_using_forall
|
|
# 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}
|