208 lines · plain
1//===- TestTilingInterfaceTransformOps.td -----------------*- tablegen -*-===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#ifndef TEST_TILINGINTERFACE_TRANSFORM_OPS10#define TEST_TILINGINTERFACE_TRANSFORM_OPS11 12include "mlir/Dialect/SCF/IR/DeviceMappingInterface.td"13include "mlir/Dialect/Transform/IR/TransformDialect.td"14include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"15include "mlir/Dialect/Transform/IR/TransformTypes.td"16include "mlir/Interfaces/SideEffectInterfaces.td"17include "mlir/IR/OpBase.td"18 19// Those operations in this file are meant for testing the tiling interface20// transformations using scf operations. Over time these testing options21// might be useful transformations in their own right. Move these over22// as transform ops in the main repo (also find a proper place for them)23 24def TestFuseAndYieldOp : Op<Transform_Dialect, "test.fuse_and_yield",25 [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,26 DeclareOpInterfaceMethods<TransformOpInterface>,27 ReportTrackingListenerFailuresOpTrait]> {28 let description = [{29 Tiles the operations pointed to by the target handle, fuses their30 producers greedily using the options provided as attributes.31 It also yields some of the fused producers for testing.32 33 On success returns the tiled operations as well as generated loops. Emits34 a definite failure if tiling fails.35 }];36 37 let arguments =38 (ins TransformHandleTypeInterface:$target,39 DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_sizes,40 DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_interchange,41 DefaultValuedAttr<BoolAttr, "false">:$use_forall);42 let results = (outs TransformHandleTypeInterface:$transfomed,43 Variadic<TransformHandleTypeInterface>:$loops);44 45 let assemblyFormat = [{46 $target ($tile_sizes^)? (`interchange` $tile_interchange^)?47 (`use_forall` $use_forall^)? attr-dict 48 `:` functional-type(operands, results)49 }];50}51 52def TestFuseConsumerUsingSliceOp : Op<Transform_Dialect, "test.fuse_consumer_using_slice",53 [AttrSizedOperandSegments,54 DeclareOpInterfaceMethods<TransformOpInterface>,55 DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,56 ReportTrackingListenerFailuresOpTrait]> {57 let description = [{58 For the `insert_slice`-like operations (that are typically generated through tiling),59 within the loop nests passed in as `loops` (that are typically generated through tiling),60 find the consumer that these slices map to (have to be the same consumer) and fuse61 the consumer into the loop.62 63 Returns a handle to the original consumer operation and the consumer operation after64 fusion.65 }];66 67 let arguments = (ins 68 Variadic<TransformHandleTypeInterface>:$targets,69 Variadic<TransformHandleTypeInterface>:$loops,70 DefaultValuedAttr<I32Attr, "1">:$num_consumer_to_fuse);71 let results = (outs TransformHandleTypeInterface:$consumer,72 TransformHandleTypeInterface:$fused_consumer);73 74 let assemblyFormat = [{75 $targets `in` `(` $loops `)`76 (`num_consumer_to_fuse` `=` $num_consumer_to_fuse^)? 77 attr-dict `:` functional-type(operands, results)78 }];79}80 81def TestFuseConsumerOp : Op<Transform_Dialect, "test.fuse_consumer",82 [DeclareOpInterfaceMethods<TransformOpInterface>,83 DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,84 ReportTrackingListenerFailuresOpTrait]> {85 let description = [{86 For the `consumer` that uses the result of the outer-most loop of a loop nest passed in 87 as `loops` (that are typically generated through tiling), fuse the consumer into the88 loop.89 90 Returns a handle to the consumer operation after fusion and the loops that might be91 modified.92 }];93 94 let arguments = (ins 95 TransformHandleTypeInterface:$consumer,96 Variadic<TransformHandleTypeInterface>:$loops);97 let results = (outs TransformHandleTypeInterface:$fused_consumer,98 Variadic<TransformHandleTypeInterface>:$result_loops);99 100 let assemblyFormat = [{101 $consumer `into` `(` $loops `)`102 attr-dict `:` functional-type(operands, results)103 }];104}105 106 107def TestTileUsingForallOp : Op<Transform_Dialect, "test.tile_using_forall",108 [DeclareOpInterfaceMethods<TransformOpInterface>,109 DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,110 ReportTrackingListenerFailuresOpTrait]> {111 let description = [{112 Test operation use to test tiling using TilingInterface and scf.forall for113 the loop constructs. This is similar to114 `transform.structured.tile_using_for`. Use of this operation is an115 intermediate state and will be replaced in due course with either116 `transform.structured.tile_using_for` or117 `transform.structured.tile_using_forall`.118 119 On success returns the tiled operations as well as generated loops. Emits120 a definite failure if tiling fails.121 }];122 123 let arguments = (ins TransformHandleTypeInterface:$target,124 DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_sizes,125 DefaultValuedOptionalAttr<I64ArrayAttr, "{}">:$interchange,126 OptionalAttr<DeviceMappingArrayAttr>:$mapping);127 let results = (outs TransformHandleTypeInterface:$tiled_op,128 Variadic<TransformHandleTypeInterface>:$loops);129 130 let assemblyFormat = [{131 $target ($tile_sizes^)? (`interchange` `=` $interchange^)?132 (`mapping` `=` $mapping^)?133 attr-dict `:` functional-type(operands, results)134 }];135}136 137def TestFuseUsingForallOp : Op<Transform_Dialect, "test.fuse_using_forall",138 [DeclareOpInterfaceMethods<TransformOpInterface>,139 DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,140 ReportTrackingListenerFailuresOpTrait]> {141 let description = [{142 Test operation to tile the operation pointed to by the target handle and143 fuses their producers greedily using the options provided as attributes.144 This operation uses scf.forall for the loop construct.145 }];146 let arguments = (ins TransformHandleTypeInterface:$root_op,147 DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_sizes,148 DefaultValuedOptionalAttr<I64ArrayAttr, "{}">:$interchange,149 OptionalAttr<DeviceMappingArrayAttr>:$mapping);150 let results = (outs TransformHandleTypeInterface:$tiled_ops,151 Variadic<TransformHandleTypeInterface>:$loops);152 153 let assemblyFormat = [{154 $root_op ($tile_sizes^)? (`interchange` $interchange^)?155 (`mapping` `=` $mapping^)?156 attr-dict `:` functional-type(operands, results)157 }];158}159 160def TestTileAndFuseOuterParallelPartialReductionOp : Op<161 Transform_Dialect, "test.tile_and_fuse_outer_parallel_partial_reduction",162 [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,163 DeclareOpInterfaceMethods<TransformOpInterface>,164 ReportTrackingListenerFailuresOpTrait]> {165 let description = [{166 Test operation to tile an operation using partial reduction with167 outer parallel strategy, and to fuse its producers.168 }];169 170 let arguments = (ins TransformHandleTypeInterface:$root_op,171 DefaultValuedAttr<I64ArrayAttr, "{}">:$reduction_dims,172 DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_sizes,173 OptionalAttr<DeviceMappingArrayAttr>:$mapping);174 175 let results = (outs TransformHandleTypeInterface:$tiled_ops,176 Variadic<TransformHandleTypeInterface>:$loops);177 let assemblyFormat = [{178 $root_op (`reduction_dims` `=` $reduction_dims^)?179 (`tile_sizes` `=` $tile_sizes^)? (`mapping` `=` $mapping^)?180 attr-dict `:` functional-type(operands, results)181 }];182}183 184def TestTileUsingCustomLoopOp : Op<185 Transform_Dialect, "test.tile_using_custom_loop",186 [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,187 DeclareOpInterfaceMethods<TransformOpInterface>,188 ReportTrackingListenerFailuresOpTrait]> {189 let description = [{190 Test Transform op to tile an operation using custom loops.191 192 The test just folds all the loops and into a single loop and then193 delinearizes the indices.194 }];195 196 let arguments = (ins TransformHandleTypeInterface:$root_op,197 DefaultValuedAttr<I64ArrayAttr, "{}">:$tile_sizes);198 let results = (outs TransformHandleTypeInterface:$tiled_ops,199 Variadic<TransformHandleTypeInterface>:$loops);200 201 let assemblyFormat = [{202 $root_op `tile_sizes` `=` $tile_sizes203 attr-dict `:` functional-type(operands, results)204 }];205}206 207#endif // TEST_TILINGINTERFACE_TRANSFORM_OPS208