Skip to content

Commit

Permalink
[mlir][tensor][linalg] Move Pack/Unpack Ops to Linalg (extra)
Browse files Browse the repository at this point in the history
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
  • Loading branch information
banach-space committed Jan 24, 2025
1 parent 55f4c72 commit f350dd2
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 3 deletions.
8 changes: 8 additions & 0 deletions mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,14 @@ def LinalgConvolutionOpInterface : OpInterface<"ConvolutionOpInterface"> {
];
}

// TODO:
def LinalgRelayoutOpInterface : OpInterface<"RelayoutOpInterface"> {
let description = [{
TODO
}];
let cppNamespace = "::mlir::linalg";
}

def LinalgFillOpInterface : OpInterface<"FillOpInterface"> {
let description = [{
A fill operation is defined in general terms:
Expand Down
3 changes: 2 additions & 1 deletion mlir/include/mlir/Dialect/Linalg/IR/LinalgRelayoutOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ include "mlir/Dialect/Linalg/IR/LinalgBase.td"
include "mlir/Interfaces/DestinationStyleOpInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
include "mlir/IR/OpAsmInterface.td"

//===----------------------------------------------------------------------===//
Expand All @@ -27,7 +28,7 @@ include "mlir/IR/OpAsmInterface.td"
class Linalg_RelayoutOp<string mnemonic, list<Trait> traits = []> :
Op<Linalg_Dialect, mnemonic, !listconcat(traits, [
DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
DestinationStyleOpInterface,
DestinationStyleOpInterface, LinalgRelayoutOpInterface,
ConditionallySpeculatable, NoMemoryEffect,
DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
TypesMatchWith<"result type matches type of dest",
Expand Down
5 changes: 3 additions & 2 deletions mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3935,9 +3935,10 @@ struct FoldTensorCastProducerOp
LogicalResult matchAndRewrite(DestinationStyleOpInterface op,
PatternRewriter &rewriter) const override {

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

SmallVector<Type> newResultTypes(op->getResultTypes());
Expand Down

0 comments on commit f350dd2

Please sign in to comment.