[mlir][linalg] Memoize indexing map generation.

[nicolasvasilache]

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

654

I expect this will show up in the printed form?
If so, can we expose the magic string as some static method and make the printer elide it?

[stellaraccident]

Address comments.

[stellaraccident]

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

654

While this is split across generators, it is hard to do generically. I just added it to the elidedAttributes in the printer explicitly for now.

[mehdi_amini]

Can you provide more context in the description?

[nicolasvasilache]

IANM this functionality has now landed as part of another revision.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

static void printNamedStructuredOpResults(OpAsmPrinter &p,
   if (resultTypes.empty())
     return;
   p.printOptionalArrowTypeList(resultTypes);
 }
 
 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.
   printCommonStructuredOpParts(p, op);
 
   // Results printing.
   printNamedStructuredOpResults(p, op.result_tensors().getTypes());
 

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

)FMT";
     // Indexing maps.
     {
       // Parameters:
       // {0}: Class name
       // {1}: Comma-separated list of dimension variable names.
       // {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";
 
       unsigned dimCount = firstMap.getNumDims();
 
       // Generate a comma-separated list of dim identifiers to be passed to
       // bindDims, ensuring tht AffineExpr identifiers are bound in the right
       // order to the proper AffineDimExpr.