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

[hanchung]

Summary

Depends On D97312

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc

 ) {
   C(b, m, n) = std_addf<k>(C(b, m, n), std_mulf(A(b, m, k), B(k, n)));
 }
 
 // Test attribute usage in affine expressions
 // 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
 // ODS-LABEL: def Test6Op
 // ODS:       let summary = [{ My magic op. }];
 // ODS-NEXT:  let description = [{
 // ODS-NEXT:    It has two inputs.
 // ODS-NEXT:    It has one output.

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

private:
 
   /// Tensors are per TC def.
   llvm::StringMap<RegisteredTensor> registeredTensors;
   unsigned nextRegisteredTensorIndex;
 
   /// 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;
 };
 } // namespace
 
 namespace llvm {
 
auto tryToParseAttrUse = [&]() -> llvm::Optional<AffineExpr> {
     auto it = registeredAttrs.find(attrName.str());
     if (it == registeredAttrs.end())
       return llvm::None;
 
     AttrUse result;
     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.
+    auto symbolIt = registeredAttrUseToSymbol.find(result.getKey());
+    if (symbolIt == registeredAttrUseToSymbol.end()) {
       result.symbol = getAffineSymbolExpr(symbols.size(), parser.context);
-    // Merely for taking the index. We don't reuse anyway.
       symbols.emplace_back("<attr-use>", result.symbol);
+      registeredAttrUseToSymbol[result.getKey()] = result.symbol;
+    } else {
+      result.symbol = symbolIt->second;
+    }
 
     attrUses.push_back(result);
 
     return result.symbol;
   };
 
   AffineParser affineParser(parser, createAffineBareId, tryToParseAttrUse, dims,
                             symbols);