20 #include "llvm/ADT/SmallVector.h"
21 #include "llvm/ADT/TypeSwitch.h"
24 #define GEN_PASS_DEF_LINALGNAMEDOPCONVERSION
25 #include "mlir/Dialect/Linalg/Passes.h.inc"
32 return llvm::to_vector<2>(llvm::seq<int64_t>(start, end));
40 auto linalgOp = dyn_cast<LinalgOp>(operation);
42 if (!linalgOp || !linalgOp.hasTensorSemantics())
47 auto kernelTy = dyn_cast<RankedTensorType>(kernel.
getType());
48 auto initTy = dyn_cast<RankedTensorType>(init.
getType());
49 auto resultTy = dyn_cast<RankedTensorType>(result.getType());
50 if (!kernelTy || !initTy || !resultTy)
53 if (kernelTy.getDimSize(3) != 1)
60 {kernelTy.getDimSize(0), kernelTy.getDimSize(1), kernelTy.getDimSize(2)},
61 kernelTy.getElementType());
62 auto collapsedKernel = rewriter.
create<tensor::CollapseShapeOp>(
63 loc, newKernelTy, kernel, collapsedKernelDims);
71 initTy.getDimSize(2), initTy.getDimSize(3)},
72 initTy.getElementType());
73 auto collapsedInit = rewriter.
create<tensor::CollapseShapeOp>(
74 loc, newInitTy, init, collapsedInitDims);
79 .Case<DepthwiseConv2DNhwcHwcmOp>([&](
auto op) {
81 return rewriter.
create<DepthwiseConv2DNhwcHwcOp>(
82 loc, newInitTy,
ValueRange{input, collapsedKernel},
85 .Case<DepthwiseConv2DNhwcHwcmQOp>([&](
auto op) {
87 return rewriter.
create<DepthwiseConv2DNhwcHwcQOp>(
88 loc, newInitTy,
ValueRange{input, collapsedKernel, iZp, kZp},
91 .Default([](
Operation *op) {
return nullptr; });
94 for (
auto attr : preservedAttrs)
95 newConv->
setAttr(attr.getName(), attr.getValue());
99 operation, resultTy, newConv->
getResult(0), collapsedInitDims);
104 struct SimplifyDepthwiseConvOp
110 Operation *operation = op.getOperation();
111 Value input = op.getDpsInputOperand(0)->get();
112 Value kernel = op.getDpsInputOperand(1)->get();
113 Value init = op.getDpsInitOperand(0)->get();
115 auto stride = op.getStrides();
116 auto dilation = op.getDilations();
119 nullptr, init, stride, dilation,
124 struct SimplifyDepthwiseConvQOp
130 Operation *operation = op.getOperation();
131 Value input = op.getDpsInputOperand(0)->get();
132 Value kernel = op.getDpsInputOperand(1)->get();
133 Value iZp = op.getDpsInputOperand(2)->get();
134 Value kZp = op.getDpsInputOperand(3)->get();
135 Value init = op.getDpsInitOperand(0)->get();
137 auto stride = op.getStrides();
138 auto dilation = op.getDilations();
141 init, stride, dilation, rewriter);
145 struct LinalgNamedOpConversionPass
146 :
public impl::LinalgNamedOpConversionBase<LinalgNamedOpConversionPass> {
147 LinalgNamedOpConversionPass() =
default;
148 LinalgNamedOpConversionPass(
const LinalgNamedOpConversionPass &) =
default;
150 void runOnOperation()
override {
155 return signalPassFailure();
162 patterns.
add<SimplifyDepthwiseConvOp, SimplifyDepthwiseConvQOp>(
167 return std::make_unique<LinalgNamedOpConversionPass>();
static llvm::SmallVector< int64_t > getIndicesVector(int start, int end)
static LogicalResult matchAndReplaceDepthwiseConv(Operation *operation, Value input, Value kernel, Value iZp, Value kZp, Value init, Attribute stride, Attribute dilation, PatternRewriter &rewriter)
Attributes are known-constant values of operations.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Operation is the basic unit of execution within MLIR.
OpResult getResult(unsigned idx)
Get the 'idx'th result of this operation.
MLIRContext * getContext()
Return the context this operation is associated with.
Location getLoc()
The source location the operation was defined or derived from.
void setAttr(StringAttr name, Attribute value)
If the an attribute exists with the specified name, change it to the new value.
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
MLIRContext * getContext() const
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replaces the result op with a new op that is created without verification.
This class provides an abstraction over the different types of ranges over Values.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Type getType() const
Return the type of this value.
void populateLinalgNamedOpConversionPatterns(RewritePatternSet &patterns)
Patterns to convert from one named op to another.
SmallVector< NamedAttribute > getPrunedAttributeList(OpTy op)
Returns an attribute list that excludes pre-defined attributes.
Include the generated interface declarations.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
LogicalResult applyPatternsAndFoldGreedily(Region ®ion, const FrozenRewritePatternSet &patterns, GreedyRewriteConfig config=GreedyRewriteConfig(), bool *changed=nullptr)
Rewrite ops in the given region, which must be isolated from above, by repeatedly applying the highes...
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
std::unique_ptr< Pass > createLinalgNamedOpConversionPass()
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
This class represents an efficient way to signal success or failure.
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...