[Mlir-commits] [mlir] 21895a2 - [mlir][linalg] Reuse the symbol if attribute uses are identical.

[Hanhan Wang]

Author: Hanhan Wang
Date: 2021-02-24T11:42:13-08:00
New Revision: 21895a2beff7fcd92441c884de7c04f324996c79

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

LOG: [mlir][linalg] Reuse the symbol if attribute uses are identical.

Depends On D97312

Reviewed By: antiagainst

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

Added: 
    

Modified: 
    mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc
    mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc b/mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc
index 7eea58869c1a..f670ac9a3c05 100644
--- a/mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc
+++ b/mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc
@@ -122,22 +122,26 @@ attr(
 // IMPL-LABEL: ArrayAttr Test5Op::indexing_maps() {
 // IMPL: auto cst0 = getAffineConstantExpr(strides().getValue<int>({ 0 }), context);
 // IMPL: auto cst1 = getAffineConstantExpr(strides().getValue<int>({ 1 }), context);
+// IMPL: auto cst2 = getAffineConstantExpr(strides().getValue<int>({ 0 }), context);
+// IMPL: auto cst3 = getAffineConstantExpr(strides().getValue<int>({ 1 }), context);
 // IMPL: auto map0 = AffineMap::get(7, 9, {d0, d1 * s7 + d4, d2 * s8 + d5, d6}, context);
-// IMPL: map0 = map0.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst0, cst1 }, 7, 0);
+// IMPL: map0 = map0.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst2, cst3 }, 7, 0);
 // IMPL: map0 = simplifyAffineMap(map0);
 // IMPL: auto map1 = AffineMap::get(7, 9, {d3, d4, d5, d6}, context);
-// IMPL: map1 = map1.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst0, cst1 }, 7, 0);
+// IMPL: map1 = map1.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst2, cst3 }, 7, 0);
 // IMPL: map1 = simplifyAffineMap(map1);
 // IMPL: auto map2 = AffineMap::get(7, 7, {d0, d1, d2, d3}, context);
-// IMPL: map2 = map2.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst0, cst1 }, 7, 0);
+// IMPL: map2 = map2.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst2, cst3 }, 7, 0);
 // IMPL: map2 = simplifyAffineMap(map2);
 // IMPL: return {{.+}}.getAffineMapArrayAttr({ map0, map1, map2 });
 //
 ods_def<Test5Op>:
 def test5(I: f32(N, H, W, C), K: f32(F, KH, KW, C)) -> (O: f32(N, H, W, F))
      attr(strides: 2xi32) {
-  O(n, h, w, f) = std_addf<kh, kw>(std_mulf(
-    I(n, h * strides[0] + kh, w * strides[1] + kw, c), K(f, kh, kw, c)));
+  O(n, h, w, f) = std_addf<kh, kw>(
+      std_mulf(std_addf(I(n, h * strides[0] + kh, w * strides[1] + kw, c),
+                        I(n, h * strides[0] + kh, w * strides[1] + kw, c)),
+               K(f, kh, kw, c)));
 }
 
 // Test documentation

diff  --git a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp
index 52fd9fbcd904..794e14780ea2 100644
--- a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp
+++ b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp
@@ -1192,6 +1192,9 @@ class TCParser {
   /// Attributes are per TC def.
   std::map<std::string, RegisteredAttr> registeredAttrs;
 
+  /// A map from AttrUse to AffineExpr symbol.
+  llvm::StringMap<AffineExpr> registeredAttrUseToSymbol;
+
   StringRef docString;
 
   Parser &parser;
@@ -1298,12 +1301,14 @@ TCParser::parseAffineExprs(EagerDiscoveryMode discoveryMode,
     if (failed(parseAttrUse(result)))
       return llvm::None;
 
-    // We create a new symbol for each attribute usage without reuse. This is
-    // fine given these symbols will be replaced with constants and folded away
-    // for concrete op instances.
-    result.symbol = getAffineSymbolExpr(symbols.size(), parser.context);
-    // Merely for taking the index. We don't reuse anyway.
-    symbols.emplace_back("<attr-use>", result.symbol);
+    auto symbolIt = registeredAttrUseToSymbol.find(result.getKey());
+    if (symbolIt == registeredAttrUseToSymbol.end()) {
+      result.symbol = getAffineSymbolExpr(symbols.size(), parser.context);
+      symbols.emplace_back("<attr-use>", result.symbol);
+      registeredAttrUseToSymbol[result.getKey()] = result.symbol;
+    } else {
+      result.symbol = symbolIt->second;
+    }
 
     attrUses.push_back(result);