diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp --- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp +++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp @@ -2333,8 +2333,11 @@ template 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. diff --git a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp --- a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp +++ b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp @@ -651,11 +651,18 @@ // {2}: Statements static const char structuredOpIndexingMapsFormat[] = R"FMT( ArrayAttr {0}::indexing_maps() { + static const char memoizeAttr[] = "linalg.memoized_indexing_maps"; + ArrayAttr cached = getOperation()->getAttrOfType(memoizeAttr); + if (cached) + return cached; + MLIRContext *context = getContext(); auto symbolBindings = getSymbolBindings(*this); SmallVector maps; {2} - return Builder(context).getAffineMapArrayAttr(maps); + cached = Builder(context).getAffineMapArrayAttr(maps); + getOperation()->setAttr(memoizeAttr, cached); + return cached; } )FMT";