Unique AffineDimExpr, AffineSymbolExpr, AffineConstantExpr, and allocate these

from the bump pointer allocator.

- delete AffineExpr destructors.

PiperOrigin-RevId: 205943807

mlir/include/mlir/IR/AffineExpr.h

@@ -72,6 +72,7 @@ public:
 
 protected:
   explicit AffineExpr(Kind kind) : kind(kind) {}
+  ~AffineExpr() {}
 
 private:
   AffineExpr(const AffineExpr&) = delete;
@@ -112,6 +113,7 @@ protected:
   AffineExpr *const rhs;
 
 private:
+  ~AffineBinaryOpExpr() = delete;
   // Simplification prior to construction of binary affine op expressions.
   static AffineExpr *simplifyAdd(AffineExpr *lhs, AffineExpr *rhs,
                                  MLIRContext *context);
@@ -143,6 +145,7 @@ public:
   }
 
 private:
+  ~AffineDimExpr() = delete;
   explicit AffineDimExpr(unsigned position)
       : AffineExpr(Kind::DimId), position(position) {}
 
@@ -167,6 +170,7 @@ public:
   }
 
 private:
+  ~AffineSymbolExpr() = delete;
   explicit AffineSymbolExpr(unsigned position)
       : AffineExpr(Kind::SymbolId), position(position) {}
 
@@ -187,6 +191,7 @@ public:
   }
 
 private:
+  ~AffineConstantExpr() = delete;
   explicit AffineConstantExpr(int64_t constant)
       : AffineExpr(Kind::Constant), constant(constant) {}
 

mlir/lib/IR/MLIRContext.cpp

@@ -200,6 +200,13 @@ public:
   DenseMap<std::tuple<unsigned, AffineExpr *, AffineExpr *>, AffineExpr *>
       affineExprs;
 
+  // Uniqui'ing of AffineDimExpr, AffineSymbolExpr's by their position.
+  std::vector<AffineDimExpr *> dimExprs;
+  std::vector<AffineSymbolExpr *> symbolExprs;
+
+  // Uniqui'ing of AffineConstantExpr using constant value as key.
+  DenseMap<int64_t, AffineConstantExpr *> constExprs;
+
   /// Integer type uniquing.
   DenseMap<unsigned, IntegerType *> integers;
 
@@ -718,21 +725,50 @@ AffineExpr *AffineBinaryOpExpr::get(AffineExpr::Kind kind, AffineExpr *lhs,
 }
 
 AffineDimExpr *AffineDimExpr::get(unsigned position, MLIRContext *context) {
-  // TODO(bondhugula): complete this
-  // FIXME: this should be POD
-  return new AffineDimExpr(position);
+  auto &impl = context->getImpl();
+
+  // Check if we need to resize.
+  if (position >= impl.dimExprs.size())
+    impl.dimExprs.resize(position + 1, nullptr);
+
+  auto *&result = impl.dimExprs[position];
+  if (result)
+    return result;
+
+  result = impl.allocator.Allocate<AffineDimExpr>();
+  // Initialize the memory using placement new.
+  new (result) AffineDimExpr(position);
+  return result;
 }
 
 AffineSymbolExpr *AffineSymbolExpr::get(unsigned position,
                                         MLIRContext *context) {
-  // TODO(bondhugula): complete this
-  // FIXME: this should be POD
-  return new AffineSymbolExpr(position);
+  auto &impl = context->getImpl();
+
+  // Check if we need to resize.
+  if (position >= impl.symbolExprs.size())
+    impl.symbolExprs.resize(position + 1, nullptr);
+
+  auto *&result = impl.symbolExprs[position];
+  if (result)
+    return result;
+
+  result = impl.allocator.Allocate<AffineSymbolExpr>();
+  // Initialize the memory using placement new.
+  new (result) AffineSymbolExpr(position);
+  return result;
 }
 
 AffineConstantExpr *AffineConstantExpr::get(int64_t constant,
                                             MLIRContext *context) {
-  // TODO(bondhugula): complete this
-  // FIXME: this should be POD
-  return new AffineConstantExpr(constant);
+  auto &impl = context->getImpl();
+  auto *&result = impl.constExprs[constant];
+
+  if (result)
+    return result;
+
+  result = impl.allocator.Allocate<AffineConstantExpr>();
+  // Initialize the memory using placement new.
+  new (result) AffineConstantExpr(constant);
+  return result;
 }

mlir/test/IR/affine-map.mlir

@@ -9,14 +9,20 @@
 // CHECK: #map{{[0-9]+}} = () -> (0)
 #map2 = () -> (0)
 
+// All three maps are unique'd as one map and so there
+// should be only one output.
 // CHECK: #map{{[0-9]+}} = (d0, d1) -> ((d0 + 1), d1)
-#map3 = (i, j) -> (i+1, j)
+#map3  = (i, j) -> (i+1, j)
+// CHECK-EMPTY
+#map3a = (i, j) -> (1+i, j)
+// CHECK-EMPTY
+#map3b = (i, j) -> (2+3-2*2+i, j)
+
+// CHECK: #map{{[0-9]+}} = (d0, d1) -> ((d0 + 2), d1)
+#map4  = (i, j) -> (3+3-2*2+i, j)
 
 // CHECK: #map{{[0-9]+}} = (d0, d1) [s0] -> ((d0 + s0), d1)
-#map4 = (i, j) [s0] -> (i + s0, j)
-
-// CHECK: #map{{[0-9]+}} = (d0, d1) -> ((d0 + 1), d1)
-#map5 = (i, j) -> (1+i, j)
+#map5 = (i, j) [s0] -> (i + s0, j)
 
 // CHECK: #map{{[0-9]+}} = (d0, d1) [s0] -> ((d0 + s0), (d1 + 5))
 #map6 = (i, j) [s0] -> (i + s0, j + 5)