[mlir][linalg] Memoize indexing map generation.

Differential Revision: https://reviews.llvm.org/D97602
This commit is contained in:
Stella Laurenzo 2021-02-26 13:11:02 -08:00
parent 7c724a896f
commit d36a15de1f
2 changed files with 13 additions and 3 deletions

View File

@ -2333,8 +2333,11 @@ static void printNamedStructuredOpResults(OpAsmPrinter &p,
template <typename NamedStructuredOpType>
static void printNamedStructuredOp(OpAsmPrinter &p, NamedStructuredOpType op) {
p << op.getOperationName();
p.printOptionalAttrDict(op->getAttrs(),
/*elidedAttrs=*/{"operand_segment_sizes"});
p.printOptionalAttrDict(
op->getAttrs(),
/*elidedAttrs=*/{"operand_segment_sizes",
// See generated code in mlir-linalg-yaml-gen.cpp
"linalg.memoized_indexing_maps"});
// Printing is shared with generic ops, except for the region and
// attributes.

View File

@ -651,11 +651,18 @@ static SmallVector<AffineExpr> getSymbolBindings({0} self) {
// {2}: Statements
static const char structuredOpIndexingMapsFormat[] = R"FMT(
ArrayAttr {0}::indexing_maps() {
static const char memoizeAttr[] = "linalg.memoized_indexing_maps";
ArrayAttr cached = getOperation()->getAttrOfType<ArrayAttr>(memoizeAttr);
if (cached)
return cached;
MLIRContext *context = getContext();
auto symbolBindings = getSymbolBindings(*this);
SmallVector<AffineMap> maps;
{2}
return Builder(context).getAffineMapArrayAttr(maps);
cached = Builder(context).getAffineMapArrayAttr(maps);
getOperation()->setAttr(memoizeAttr, cached);
return cached;
}
)FMT";