a9efcbf490
This patch enables continuous tiling of a target structured op using diminishing tile sizes. In cases where the tensor dimensions are not exactly divisible by the tile size, we are left with leftover tensor chunks that are irregularly tiled. This approach enables tiling of the leftover chunk with a smaller tile size and repeats this process recursively using exponentially diminishing tile sizes. This eventually generates a chain of loops that apply tiling using diminishing tile sizes. Adds `continuous_tile_sizes` op to the transform dialect. This op, when given a tile size and a dimension, computes a series of diminishing tile sizes that can be used to tile the target along the given dimension. Additionally, this op also generates a series of chunk sizes that the corresponding tile sizes should be applied to along the given dimension. Adds `multiway` attribute to `transform.structured.split` that enables multiway splitting of a single target op along the given dimension, as specified in a list enumerating the chunk sizes.
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, chunk_sizes=42)
|
|
structured.SplitOp(split.results[0], dimension=3, chunk_sizes=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}
|