Skip to content

Commit f350dd2

Browse files
committed
[mlir][tensor][linalg] Move Pack/Unpack Ops to Linalg (extra)
This is merely moving code around, no new functionality is added. PATCH 5: Create `LinalgRelayoutOpInterface` to be able to exclude `linalg::PackOp` + `linalg::UnpackOp` Ops from patterns/folders outside the Linalg dialect, e.g. `FoldTensorCastProducerOp` from the Tensor dialect. Note that there's `FoldTensorCastUnPackOp` and `FoldTensorCastPackOp` in LinalgOps.cpp (i.e. Linalg dialect) that provides similar folder (but which fold "correctly"). See e.g. #121393 and #114559 for context. CONTEXT: This change was discussed in the following RFC: * https://discourse.llvm.org/t/rfc-move-tensor-pack-and-tensor-unpack-into-linalg
1 parent 55f4c72 commit f350dd2

File tree

3 files changed

+13
-3
lines changed

3 files changed

+13
-3
lines changed

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,14 @@ def LinalgConvolutionOpInterface : OpInterface<"ConvolutionOpInterface"> {
178178
];
179179
}
180180

181+
// TODO:
182+
def LinalgRelayoutOpInterface : OpInterface<"RelayoutOpInterface"> {
183+
let description = [{
184+
TODO
185+
}];
186+
let cppNamespace = "::mlir::linalg";
187+
}
188+
181189
def LinalgFillOpInterface : OpInterface<"FillOpInterface"> {
182190
let description = [{
183191
A fill operation is defined in general terms:

mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ include "mlir/Dialect/Linalg/IR/LinalgBase.td"
1818
include "mlir/Interfaces/DestinationStyleOpInterface.td"
1919
include "mlir/Interfaces/SideEffectInterfaces.td"
2020
include "mlir/Interfaces/InferTypeOpInterface.td"
21+
include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
2122
include "mlir/IR/OpAsmInterface.td"
2223

2324
//===----------------------------------------------------------------------===//
@@ -27,7 +28,7 @@ include "mlir/IR/OpAsmInterface.td"
2728
class Linalg_RelayoutOp<string mnemonic, list<Trait> traits = []> :
2829
Op<Linalg_Dialect, mnemonic, !listconcat(traits, [
2930
DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
30-
DestinationStyleOpInterface,
31+
DestinationStyleOpInterface, LinalgRelayoutOpInterface,
3132
ConditionallySpeculatable, NoMemoryEffect,
3233
DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
3334
TypesMatchWith<"result type matches type of dest",

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3935,9 +3935,10 @@ struct FoldTensorCastProducerOp
39353935
LogicalResult matchAndRewrite(DestinationStyleOpInterface op,
39363936
PatternRewriter &rewriter) const override {
39373937

3938-
// Reject PackOp/UnpackOp - there are dedicated patterns for that instead.
3938+
// Reject PackOp/UnpackOp (i.e. RelayoutOps) - there are dedicated patterns
3939+
// for that instead.
39393940
if (!foldTensorCastPrecondition(op) ||
3940-
isa<linalg::PackOp, linalg::UnPackOp>(*op))
3941+
isa<linalg::RelayoutOpInterface>(*op))
39413942
return failure();
39423943

39433944
SmallVector<Type> newResultTypes(op->getResultTypes());

0 commit comments

Comments
 (0)