diff --git a/mlir/include/mlir/Analysis/Presburger/PWMAFunction.h b/mlir/include/mlir/Analysis/Presburger/PWMAFunction.h
--- a/mlir/include/mlir/Analysis/Presburger/PWMAFunction.h
+++ b/mlir/include/mlir/Analysis/Presburger/PWMAFunction.h
@@ -134,22 +134,26 @@
 ///
 /// Support is provided to compare equality of two such functions as well as
 /// finding the value of the function at a point.
-class PWMAFunction : public PresburgerSpace {
+class PWMAFunction {
 public:
   PWMAFunction(const PresburgerSpace &space, unsigned numOutputs)
-      : PresburgerSpace(space), numOutputs(numOutputs) {
-    assert(getNumDomainIds() == 0 && "Set type space should zero domain ids.");
-    assert(getNumLocalIds() == 0 && "PWMAFunction cannot have local ids.");
+      : space(space), numOutputs(numOutputs) {
+    assert(space.getNumDomainIds() == 0 &&
+           "Set type space should zero domain ids.");
+    assert(space.getNumLocalIds() == 0 &&
+           "PWMAFunction cannot have local ids.");
     assert(numOutputs >= 1 && "The function must output something!");
   }
 
+  const PresburgerSpace &getSpace() const { return space; }
+
   void addPiece(const MultiAffineFunction &piece);
   void addPiece(const IntegerPolyhedron &domain, const Matrix &output);
 
   const MultiAffineFunction &getPiece(unsigned i) const { return pieces[i]; }
   unsigned getNumPieces() const { return pieces.size(); }
   unsigned getNumOutputs() const { return numOutputs; }
-  unsigned getNumInputs() const { return getNumIds(); }
+  unsigned getNumInputs() const { return space.getNumIds(); }
   MultiAffineFunction &getPiece(unsigned i) { return pieces[i]; }
 
   /// Return the domain of this piece-wise MultiAffineFunction. This is the
@@ -169,6 +173,8 @@
   void dump() const;
 
 private:
+  PresburgerSpace space;
+
   /// The list of pieces in this piece-wise MultiAffineFunction.
   SmallVector<MultiAffineFunction, 4> pieces;
 
diff --git a/mlir/lib/Analysis/Presburger/PWMAFunction.cpp b/mlir/lib/Analysis/Presburger/PWMAFunction.cpp
--- a/mlir/lib/Analysis/Presburger/PWMAFunction.cpp
+++ b/mlir/lib/Analysis/Presburger/PWMAFunction.cpp
@@ -149,7 +149,7 @@
 /// Two PWMAFunctions are equal if they have the same dimensionalities,
 /// the same domain, and take the same value at every point in the domain.
 bool PWMAFunction::isEqual(const PWMAFunction &other) const {
-  if (!isSpaceCompatible(other))
+  if (!space.isSpaceCompatible(other.getSpace()))
     return false;
 
   if (!this->getDomain().isEqual(other.getDomain()))
@@ -167,7 +167,7 @@
 }
 
 void PWMAFunction::addPiece(const MultiAffineFunction &piece) {
-  assert(piece.isSpaceCompatible(*this) &&
+  assert(piece.isSpaceCompatible(space) &&
          "Piece to be added is not compatible with this PWMAFunction!");
   assert(piece.isConsistent() && "Piece is internally inconsistent!");
   assert(this->getDomain()