[llvm-branch-commits] [mlir] d36a15d - [mlir][linalg] Memoize indexing map generation.

[Stella Laurenzo via llvm-branch-commits]

Author: Stella Laurenzo
Date: 2021-03-01T21:15:40-08:00
New Revision: d36a15de1ff4d24e772233406d602c5f0b370f54

URL: https://github.com/llvm/llvm-project/commit/d36a15de1ff4d24e772233406d602c5f0b370f54
DIFF: https://github.com/llvm/llvm-project/commit/d36a15de1ff4d24e772233406d602c5f0b370f54.diff

LOG: [mlir][linalg] Memoize indexing map generation.

Differential Revision: https://reviews.llvm.org/D97602

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
    mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index acc8ff1807c1..46e5780e151f 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -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.

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
index 5578ff52d477..1dddc57f25d3 100644
--- 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 @@ 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";