[mlir][python] improve usability of Python affine construct bindings

[ftynse]

Summary

• Provide the operator overloads for constructing (semi-)affine expressions in Python by combining existing expressions with constants.
• Make AffineExpr, AffineMap and IntegerSet hashable in Python.
• Expose the AffineExpr composition functionality.

[gysit]

LGTM!

mlir/lib/Bindings/Python/IRAffine.cpp

237

Shouldn't this be an PyAffineMulExpr? It probably doe not matter since both are affine expressions internally? The same would apply below to the modulo / div operators.

587

nit: modulo

[ftynse]

Address review.

mlir/include/mlir-c/AffineExpr.h

struct name {                                                                \
     storage *ptr;                                                              \
   };                                                                           \
   typedef struct name name
 
 DEFINE_C_API_STRUCT(MlirAffineExpr, const void);
 
 #undef DEFINE_C_API_STRUCT
 
+struct MlirAffineMap;
+
 /// Gets the context that owns the affine expression.
 MLIR_CAPI_EXPORTED MlirContext
 mlirAffineExprGetContext(MlirAffineExpr affineExpr);
 
 /// Returns `true` if the two affine expressions are equal.
 MLIR_CAPI_EXPORTED bool mlirAffineExprEqual(MlirAffineExpr lhs,
                                             MlirAffineExpr rhs);
 

 MLIR_CAPI_EXPORTED bool mlirAffineExprIsMultipleOf(MlirAffineExpr affineExpr,
                                                    int64_t factor);
 
 /// Checks whether the given affine expression involves AffineDimExpr
 /// 'position'.
 MLIR_CAPI_EXPORTED bool mlirAffineExprIsFunctionOfDim(MlirAffineExpr affineExpr,
                                                       intptr_t position);
 
+/// Composes the given map with the given expression.
+MLIR_CAPI_EXPORTED MlirAffineExpr mlirAffineExprCompose(
+    MlirAffineExpr affineExpr, struct MlirAffineMap affineMap);
+
 //===----------------------------------------------------------------------===//
 // Affine Dimension Expression.
 //===----------------------------------------------------------------------===//
 
 /// Checks whether the given affine expression is a dimension expression.
 MLIR_CAPI_EXPORTED bool mlirAffineExprIsADim(MlirAffineExpr affineExpr);
 
 /// Creates an affine dimension expression with 'position' in the context.

mlir/lib/Bindings/Python/IRAffine.cpp

public:
   static constexpr const char *pyClassName = "AffineAddExpr";
   using PyConcreteAffineExpr::PyConcreteAffineExpr;
 
   static PyAffineAddExpr get(PyAffineExpr lhs, PyAffineExpr rhs) {
     MlirAffineExpr expr = mlirAffineAddExprGet(lhs, rhs);
     return PyAffineAddExpr(lhs.getContext(), expr);
   }
 
+  static PyAffineAddExpr getRHSConstant(PyAffineExpr lhs, intptr_t rhs) {
+    MlirAffineExpr expr = mlirAffineAddExprGet(
+        lhs, mlirAffineConstantExprGet(mlirAffineExprGetContext(lhs), rhs));
+    return PyAffineAddExpr(lhs.getContext(), expr);
+  }
+
+  static PyAffineAddExpr getLHSConstant(intptr_t lhs, PyAffineExpr rhs) {
+    MlirAffineExpr expr = mlirAffineAddExprGet(
+        mlirAffineConstantExprGet(mlirAffineExprGetContext(rhs), lhs), rhs);
+    return PyAffineAddExpr(rhs.getContext(), expr);
+  }
+
   static void bindDerived(ClassTy &c) {
     c.def_static("get", &PyAffineAddExpr::get);
   }
 };
 
 class PyAffineMulExpr
     : public PyConcreteAffineExpr<PyAffineMulExpr, PyAffineBinaryExpr> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAffineExprIsAMul;
   static constexpr const char *pyClassName = "AffineMulExpr";
   using PyConcreteAffineExpr::PyConcreteAffineExpr;
 
   static PyAffineMulExpr get(PyAffineExpr lhs, PyAffineExpr rhs) {
     MlirAffineExpr expr = mlirAffineMulExprGet(lhs, rhs);
     return PyAffineMulExpr(lhs.getContext(), expr);
   }
 
+  static PyAffineMulExpr getRHSConstant(PyAffineExpr lhs, intptr_t rhs) {
+    MlirAffineExpr expr = mlirAffineMulExprGet(
+        lhs, mlirAffineConstantExprGet(mlirAffineExprGetContext(lhs), rhs));
+    return PyAffineMulExpr(lhs.getContext(), expr);
+  }
+
+  static PyAffineMulExpr getLHSConstant(intptr_t lhs, PyAffineExpr rhs) {
+    MlirAffineExpr expr = mlirAffineMulExprGet(
+        mlirAffineConstantExprGet(mlirAffineExprGetContext(rhs), lhs), rhs);
+    return PyAffineMulExpr(rhs.getContext(), expr);
+  }
+
   static void bindDerived(ClassTy &c) {
     c.def_static("get", &PyAffineMulExpr::get);
   }
 };
 
 class PyAffineModExpr
     : public PyConcreteAffineExpr<PyAffineModExpr, PyAffineBinaryExpr> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAffineExprIsAMod;
   static constexpr const char *pyClassName = "AffineModExpr";
   using PyConcreteAffineExpr::PyConcreteAffineExpr;
 
   static PyAffineModExpr get(PyAffineExpr lhs, PyAffineExpr rhs) {
     MlirAffineExpr expr = mlirAffineModExprGet(lhs, rhs);
     return PyAffineModExpr(lhs.getContext(), expr);
   }
 
+  static PyAffineModExpr getRHSConstant(PyAffineExpr lhs, intptr_t rhs) {
+    MlirAffineExpr expr = mlirAffineModExprGet(
+        lhs, mlirAffineConstantExprGet(mlirAffineExprGetContext(lhs), rhs));
+    return PyAffineModExpr(lhs.getContext(), expr);
+  }
+
+  static PyAffineModExpr getLHSConstant(intptr_t lhs, PyAffineExpr rhs) {
+    MlirAffineExpr expr = mlirAffineModExprGet(
+        mlirAffineConstantExprGet(mlirAffineExprGetContext(rhs), lhs), rhs);
+    return PyAffineModExpr(rhs.getContext(), expr);
+  }
+
   static void bindDerived(ClassTy &c) {
     c.def_static("get", &PyAffineModExpr::get);
   }
 };
 
 class PyAffineFloorDivExpr
     : public PyConcreteAffineExpr<PyAffineFloorDivExpr, PyAffineBinaryExpr> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAffineExprIsAFloorDiv;
   static constexpr const char *pyClassName = "AffineFloorDivExpr";
   using PyConcreteAffineExpr::PyConcreteAffineExpr;
 
   static PyAffineFloorDivExpr get(PyAffineExpr lhs, PyAffineExpr rhs) {
     MlirAffineExpr expr = mlirAffineFloorDivExprGet(lhs, rhs);
     return PyAffineFloorDivExpr(lhs.getContext(), expr);
   }
 
+  static PyAffineFloorDivExpr getRHSConstant(PyAffineExpr lhs, intptr_t rhs) {
+    MlirAffineExpr expr = mlirAffineFloorDivExprGet(
+        lhs, mlirAffineConstantExprGet(mlirAffineExprGetContext(lhs), rhs));
+    return PyAffineFloorDivExpr(lhs.getContext(), expr);
+  }
+
+  static PyAffineFloorDivExpr getLHSConstant(intptr_t lhs, PyAffineExpr rhs) {
+    MlirAffineExpr expr = mlirAffineFloorDivExprGet(
+        mlirAffineConstantExprGet(mlirAffineExprGetContext(rhs), lhs), rhs);
+    return PyAffineFloorDivExpr(rhs.getContext(), expr);
+  }
+
   static void bindDerived(ClassTy &c) {
     c.def_static("get", &PyAffineFloorDivExpr::get);
   }
 };
 
 class PyAffineCeilDivExpr
     : public PyConcreteAffineExpr<PyAffineCeilDivExpr, PyAffineBinaryExpr> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAffineExprIsACeilDiv;
   static constexpr const char *pyClassName = "AffineCeilDivExpr";
   using PyConcreteAffineExpr::PyConcreteAffineExpr;
 
   static PyAffineCeilDivExpr get(PyAffineExpr lhs, PyAffineExpr rhs) {
     MlirAffineExpr expr = mlirAffineCeilDivExprGet(lhs, rhs);
     return PyAffineCeilDivExpr(lhs.getContext(), expr);
   }
 
+  static PyAffineCeilDivExpr getRHSConstant(PyAffineExpr lhs, intptr_t rhs) {
+    MlirAffineExpr expr = mlirAffineCeilDivExprGet(
+        lhs, mlirAffineConstantExprGet(mlirAffineExprGetContext(lhs), rhs));
+    return PyAffineCeilDivExpr(lhs.getContext(), expr);
+  }
+
+  static PyAffineCeilDivExpr getLHSConstant(intptr_t lhs, PyAffineExpr rhs) {
+    MlirAffineExpr expr = mlirAffineCeilDivExprGet(
+        mlirAffineConstantExprGet(mlirAffineExprGetContext(rhs), lhs), rhs);
+    return PyAffineCeilDivExpr(rhs.getContext(), expr);
+  }
+
   static void bindDerived(ClassTy &c) {
     c.def_static("get", &PyAffineCeilDivExpr::get);
   }
 };
 
 } // namespace
 
 bool PyAffineExpr::operator==(const PyAffineExpr &other) {

 void mlir::python::populateIRAffine(py::module &m) {
   //----------------------------------------------------------------------------
   // Mapping of PyAffineExpr and derived classes.
   //----------------------------------------------------------------------------
   py::class_<PyAffineExpr>(m, "AffineExpr", py::module_local())
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyAffineExpr::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyAffineExpr::createFromCapsule)
-      .def("__add__",
-           [](PyAffineExpr &self, PyAffineExpr &other) {
-             return PyAffineAddExpr::get(self, other);
-           })
-      .def("__mul__",
-           [](PyAffineExpr &self, PyAffineExpr &other) {
-             return PyAffineMulExpr::get(self, other);
-           })
-      .def("__mod__",
-           [](PyAffineExpr &self, PyAffineExpr &other) {
-             return PyAffineModExpr::get(self, other);
+      .def("__add__", &PyAffineAddExpr::get)
+      .def("__add__", &PyAffineAddExpr::getRHSConstant)
+      .def("__radd__", &PyAffineAddExpr::getRHSConstant)
+      .def("__mul__", &PyAffineMulExpr::get)
+      .def("__mul__", &PyAffineMulExpr::getRHSConstant)
+      .def("__rmul__", &PyAffineMulExpr::getRHSConstant)
+      .def("__mod__", &PyAffineModExpr::get)
+      .def("__mod__", &PyAffineModExpr::getRHSConstant)
+      .def("__rmod__",
+           [](PyAffineExpr &self, intptr_t other) {
+             return PyAffineModExpr::get(
+                 PyAffineConstantExpr::get(other, *self.getContext().get()),
+                 self);
            })
       .def("__sub__",
            [](PyAffineExpr &self, PyAffineExpr &other) {
              auto negOne =
                  PyAffineConstantExpr::get(-1, *self.getContext().get());
              return PyAffineAddExpr::get(self,
                                          PyAffineMulExpr::get(negOne, other));
            })
+      .def("__sub__",
+           [](PyAffineExpr &self, intptr_t other) {
+             return PyAffineAddExpr::get(
+                 self,
+                 PyAffineConstantExpr::get(-other, *self.getContext().get()));
+           })
+      .def("__rsub__",
+           [](PyAffineExpr &self, intptr_t other) {
+             return PyAffineAddExpr::getLHSConstant(
+                 other, PyAffineMulExpr::getLHSConstant(-1, self));
+           })
       .def("__eq__", [](PyAffineExpr &self,
                         PyAffineExpr &other) { return self == other; })
       .def("__eq__",
            [](PyAffineExpr &self, py::object &other) { return false; })
       .def("__str__",
            [](PyAffineExpr &self) {
              PyPrintAccumulator printAccum;
              mlirAffineExprPrint(self, printAccum.getCallback(),
                                  printAccum.getUserData());
              return printAccum.join();
            })
       .def("__repr__",
            [](PyAffineExpr &self) {
              PyPrintAccumulator printAccum;
              printAccum.parts.append("AffineExpr(");
              mlirAffineExprPrint(self, printAccum.getCallback(),
                                  printAccum.getUserData());
              printAccum.parts.append(")");
              return printAccum.join();
            })
+      .def("__hash__",
+           [](PyAffineExpr &self) {
+             return static_cast<size_t>(llvm::hash_value(self.get().ptr));
+           })
       .def_property_readonly(
           "context",
           [](PyAffineExpr &self) { return self.getContext().getObject(); })
+      .def("compose",
+           [](PyAffineExpr &self, PyAffineMap &other) {
+             return PyAffineExpr(self.getContext(),
+                                 mlirAffineExprCompose(self, other));
+           })
       .def_static(
           "get_add", &PyAffineAddExpr::get,
           "Gets an affine expression containing a sum of two expressions.")
+      .def_static("get_add", &PyAffineAddExpr::getLHSConstant,
+                  "Gets an affine expression containing a sum of a constant "
+                  "and another expression.")
+      .def_static("get_add", &PyAffineAddExpr::getRHSConstant,
+                  "Gets an affine expression containing a sum of an expression "
+                  "and a constant.")
       .def_static(
           "get_mul", &PyAffineMulExpr::get,
           "Gets an affine expression containing a product of two expressions.")
+      .def_static("get_mul", &PyAffineMulExpr::getLHSConstant,
+                  "Gets an affine expression containing a product of a "
+                  "constant and another expression.")
+      .def_static("get_mul", &PyAffineMulExpr::getRHSConstant,
+                  "Gets an affine expression containing a product of an "
+                  "expression and a constant.")
       .def_static("get_mod", &PyAffineModExpr::get,
                   "Gets an affine expression containing the modulo of dividing "
                   "one expression by another.")
+      .def_static("get_mod", &PyAffineModExpr::getLHSConstant,
+                  "Gets a semi-affine expression containing the modulo of "
+                  "dividing a constant by an expression.")
+      .def_static("get_mod", &PyAffineModExpr::getRHSConstant,
+                  "Gets an affine expression containing the module of dividing"
+                  "an expression by a constant.")
       .def_static("get_floor_div", &PyAffineFloorDivExpr::get,
                   "Gets an affine expression containing the rounded-down "
                   "result of dividing one expression by another.")
+      .def_static("get_floor_div", &PyAffineFloorDivExpr::getLHSConstant,
+                  "Gets a semi-affine expression containing the rounded-down "
+                  "result of dividing a constant by an expression.")
+      .def_static("get_floor_div", &PyAffineFloorDivExpr::getRHSConstant,
+                  "Gets an affine expression containing the rounded-down "
+                  "result of dividing an expression by a constant.")
       .def_static("get_ceil_div", &PyAffineCeilDivExpr::get,
                   "Gets an affine expression containing the rounded-up result "
                   "of dividing one expression by another.")
+      .def_static("get_ceil_div", &PyAffineCeilDivExpr::getLHSConstant,
+                  "Gets a semi-affine expression containing the rounded-up "
+                  "result of dividing a constant by an expression.")
+      .def_static("get_ceil_div", &PyAffineCeilDivExpr::getRHSConstant,
+                  "Gets an affine expression containing the rounded-up result "
+                  "of dividing an expression by a constant.")
       .def_static("get_constant", &PyAffineConstantExpr::get, py::arg("value"),
                   py::arg("context") = py::none(),
                   "Gets a constant affine expression with the given value.")
       .def_static(
           "get_dim", &PyAffineDimExpr::get, py::arg("position"),
           py::arg("context") = py::none(),
           "Gets an affine expression of a dimension at the given position.")
       .def_static(
py::class_<PyAffineMap>(m, "AffineMap", py::module_local())
            [](PyAffineMap &self) {
              PyPrintAccumulator printAccum;
              printAccum.parts.append("AffineMap(");
              mlirAffineMapPrint(self, printAccum.getCallback(),
                                 printAccum.getUserData());
              printAccum.parts.append(")");
              return printAccum.join();
            })
+      .def("__hash__",
+           [](PyAffineMap &self) {
+             return static_cast<size_t>(llvm::hash_value(self.get().ptr));
+           })
       .def_static("compress_unused_symbols",
                   [](py::list affineMaps, DefaultingPyMlirContext context) {
                     SmallVector<MlirAffineMap> maps;
                     pyListToVector<PyAffineMap, MlirAffineMap>(
                         affineMaps, maps, "attempting to create an AffineMap");
                     std::vector<MlirAffineMap> compressed(affineMaps.size());
                     auto populate = [](void *result, intptr_t idx,
                                        MlirAffineMap m) {
py::class_<PyIntegerSet>(m, "IntegerSet", py::module_local())
            [](PyIntegerSet &self) {
              PyPrintAccumulator printAccum;
              printAccum.parts.append("IntegerSet(");
              mlirIntegerSetPrint(self, printAccum.getCallback(),
                                  printAccum.getUserData());
              printAccum.parts.append(")");
              return printAccum.join();
            })
+      .def("__hash__",
+           [](PyIntegerSet &self) {
+             return static_cast<size_t>(llvm::hash_value(self.get().ptr));
+           })
       .def_property_readonly(
           "context",
           [](PyIntegerSet &self) { return self.getContext().getObject(); })
       .def(
           "dump", [](PyIntegerSet &self) { mlirIntegerSetDump(self); },
           kDumpDocstring)
       .def_static(
           "get",

mlir/lib/CAPI/IR/AffineExpr.cpp

bool mlirAffineExprIsMultipleOf(MlirAffineExpr affineExpr, int64_t factor) {
   return unwrap(affineExpr).isMultipleOf(factor);
 }
 
 bool mlirAffineExprIsFunctionOfDim(MlirAffineExpr affineExpr,
                                    intptr_t position) {
   return unwrap(affineExpr).isFunctionOfDim(position);
 }
 
+MlirAffineExpr mlirAffineExprCompose(MlirAffineExpr affineExpr,
+                                     MlirAffineMap affineMap) {
+  return wrap(unwrap(affineExpr).compose(unwrap(affineMap)));
+}
+
 //===----------------------------------------------------------------------===//
 // Affine Dimension Expression.
 //===----------------------------------------------------------------------===//
 
 bool mlirAffineExprIsADim(MlirAffineExpr affineExpr) {
   return unwrap(affineExpr).isa<AffineDimExpr>();
 }
 

mlir/test/CAPI/ir.c

if (mlirAffineMapGetNumResults(map) != 2)
     return 1;
 
   if (!mlirAffineExprEqual(mlirAffineMapGetResult(map, 0), affineDimExpr))
     return 2;
 
   if (!mlirAffineExprEqual(mlirAffineMapGetResult(map, 1), affineSymbolExpr))
     return 3;
 
+  MlirAffineExpr affineDim2Expr = mlirAffineDimExprGet(ctx, 1);
+  MlirAffineExpr composed = mlirAffineExprCompose(affineDim2Expr, map);
+  // CHECK: s1
+  mlirAffineExprDump(composed);
+  if (!mlirAffineExprEqual(composed, affineSymbolExpr))
+    return 4;
+
   return 0;
 }
 
 int printIntegerSet(MlirContext ctx) {
   MlirIntegerSet emptySet = mlirIntegerSetEmptyGet(ctx, 2, 1);
 
   // CHECK-LABEL: @printIntegerSet
   fprintf(stderr, "@printIntegerSet");

mlir/test/python/ir/affine_expr.py

with Context():
     d12 = AffineExpr.get_add(d1, d2)
     # CHECK: d1 + d2
     print(d12)
 
     d12op = d1 + d2
     # CHECK: d1 + d2
     print(d12op)
 
+    d1cst_op = d1 + 2
+    # CHECK: d1 + 2
+    print(d1cst_op)
+
+    d1cst_op2 = 2 + d1
+    # CHECK: d1 + 2
+    print(d1cst_op2)
+
     assert d12 == d12op
     assert d12.lhs == d1
     assert d12.rhs == d2
 
 
 # CHECK-LABEL: TEST: testAffineMulExpr
 @run
 def testAffineMulExpr():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     expr = AffineExpr.get_mul(d1, c2)
     # CHECK: d1 * 2
     print(expr)
 
     # CHECK: d1 * 2
     op = d1 * c2
     print(op)
 
+    # CHECK: d1 * 2
+    op_cst = d1 * 2
+    print(op_cst)
+
+    # CHECK: d1 * 2
+    op_cst2 = 2 * d1
+    print(op_cst2)
+
     assert expr == op
+    assert expr == op_cst
     assert expr.lhs == d1
     assert expr.rhs == c2
 
 
 # CHECK-LABEL: TEST: testAffineModExpr
 @run
 def testAffineModExpr():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     expr = AffineExpr.get_mod(d1, c2)
     # CHECK: d1 mod 2
     print(expr)
 
     # CHECK: d1 mod 2
     op = d1 % c2
     print(op)
 
+    # CHECK: d1 mod 2
+    op_cst = d1 % 2
+    print(op_cst)
+
+    # CHECK: 2 mod d1
+    print(2 % d1)
+
     assert expr == op
+    assert expr == op_cst
     assert expr.lhs == d1
     assert expr.rhs == c2
 
+    expr2 = AffineExpr.get_mod(c2, d1)
+    expr3 = AffineExpr.get_mod(2, d1)
+    expr4 = AffineExpr.get_mod(d1, 2)
+
+    # CHECK: 2 mod d1
+    print(expr2)
+    # CHECK: 2 mod d1
+    print(expr3)
+    # CHECK: d1 mod 2
+    print(expr4)
+
+    assert expr2 == expr3
+    assert expr4 == expr
+
 
 # CHECK-LABEL: TEST: testAffineFloorDivExpr
 @run
 def testAffineFloorDivExpr():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     expr = AffineExpr.get_floor_div(d1, c2)
     # CHECK: d1 floordiv 2
     print(expr)
 
     assert expr.lhs == d1
     assert expr.rhs == c2
 
+    expr2 = AffineExpr.get_floor_div(c2, d1)
+    expr3 = AffineExpr.get_floor_div(2, d1)
+    expr4 = AffineExpr.get_floor_div(d1, 2)
+
+    # CHECK: 2 floordiv d1
+    print(expr2)
+    # CHECK: 2 floordiv d1
+    print(expr3)
+    # CHECK: d1 floordiv 2
+    print(expr4)
+
+    assert expr2 == expr3
+    assert expr4 == expr
+
 
 # CHECK-LABEL: TEST: testAffineCeilDivExpr
 @run
 def testAffineCeilDivExpr():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     expr = AffineExpr.get_ceil_div(d1, c2)
     # CHECK: d1 ceildiv 2
     print(expr)
 
     assert expr.lhs == d1
     assert expr.rhs == c2
 
+    expr2 = AffineExpr.get_ceil_div(c2, d1)
+    expr3 = AffineExpr.get_ceil_div(2, d1)
+    expr4 = AffineExpr.get_ceil_div(d1, 2)
+
+    # CHECK: 2 ceildiv d1
+    print(expr2)
+    # CHECK: 2 ceildiv d1
+    print(expr3)
+    # CHECK: d1 ceildiv 2
+    print(expr4)
+
+    assert expr2 == expr3
+    assert expr4 == expr
+
 
 # CHECK-LABEL: TEST: testAffineExprSub
 @run
 def testAffineExprSub():
   with Context():
     d1 = AffineDimExpr.get(1)
     d2 = AffineDimExpr.get(2)
     expr = d1 - d2
     # CHECK: d1 - d2
     print(expr)
 
     assert expr.lhs == d1
     rhs = AffineMulExpr(expr.rhs)
     # CHECK: d2
     print(rhs.lhs)
     # CHECK: -1
     print(rhs.rhs)
 
+    # CHECK: d1 - 42
+    print(d1 - 42)
+    # CHECK: -d1 + 42
+    print(42 - d1)
+
+    c42 = AffineConstantExpr.get(42)
+    assert d1 - 42 == d1 - c42
+    assert 42 - d1 == c42 - d1
+
 # CHECK-LABEL: TEST: testClassHierarchy
 @run
 def testClassHierarchy():
   with Context():
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     add = AffineAddExpr.get(d1, c2)
     mul = AffineMulExpr.get(d1, c2)
with Context():
     # CHECK: True
     print(AffineAddExpr.isinstance(add))
     # CHECK: False
     print(AffineMulExpr.isinstance(add))
     # CHECK: True
     print(AffineMulExpr.isinstance(mul))
     # CHECK: False
     print(AffineAddExpr.isinstance(mul))
+
+
+# CHECK-LABEL: TEST: testCompose
+@run
+def testCompose():
+  with Context():
+    # d0 + d2.
+    expr = AffineAddExpr.get(AffineDimExpr.get(0), AffineDimExpr.get(2))
+
+    # (d0, d1, d2)[s0, s1] -> (d0 + s1, d1 + s0, d0 + d1 + d2)
+    map1 = AffineAddExpr.get(AffineDimExpr.get(0), AffineSymbolExpr.get(1))
+    map2 = AffineAddExpr.get(AffineDimExpr.get(1), AffineSymbolExpr.get(0))
+    map3 = AffineAddExpr.get(
+        AffineAddExpr.get(AffineDimExpr.get(0), AffineDimExpr.get(1)),
+        AffineDimExpr.get(2))
+    map = AffineMap.get(3, 2, [map1, map2, map3])
+
+    # CHECK: d0 + s1 + d0 + d1 + d2
+    print(expr.compose(map))
+
+
+# CHECK-LABEL: TEST: testHash
+@run
+def testHash():
+  with Context():
+    d0 = AffineDimExpr.get(0)
+    s1 = AffineSymbolExpr.get(1)
+    assert hash(d0) == hash(AffineDimExpr.get(0))
+    assert hash(d0 + s1) == hash(AffineAddExpr.get(d0, s1))
+
+    dictionary = dict()
+    dictionary[d0] = 0
+    dictionary[s1] = 1
+    assert d0 in dictionary
+    assert s1 in dictionary

mlir/test/python/ir/affine_map.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
+  return f
 
 
 # CHECK-LABEL: TEST: testAffineMapCapsule
+@run
 def testAffineMapCapsule():
   with Context() as ctx:
     am1 = AffineMap.get_empty(ctx)
   # CHECK: mlir.ir.AffineMap._CAPIPtr
   affine_map_capsule = am1._CAPIPtr
   print(affine_map_capsule)
   am2 = AffineMap._CAPICreate(affine_map_capsule)
   assert am2 == am1
   assert am2.context is ctx
 
 
-run(testAffineMapCapsule)
-
-
 # CHECK-LABEL: TEST: testAffineMapGet
+@run
 def testAffineMapGet():
   with Context() as ctx:
     d0 = AffineDimExpr.get(0)
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
 
     # CHECK: (d0, d1)[s0, s1, s2] -> ()
     map0 = AffineMap.get(2, 3, [])
with Context() as ctx:
 
     try:
       map3.get_major_submap(42)
     except ValueError as e:
       # CHECK: number of results out of bounds
       print(e)
 
 
-run(testAffineMapGet)
-
-
 # CHECK-LABEL: TEST: testAffineMapDerive
+@run
 def testAffineMapDerive():
   with Context() as ctx:
     map5 = AffineMap.get_identity(5)
 
     # CHECK: (d0, d1, d2, d3, d4) -> (d1, d2, d3)
     map123 = map5.get_submap([1, 2, 3])
     print(map123)
 
     # CHECK: (d0, d1, d2, d3, d4) -> (d0, d1)
     map01 = map5.get_major_submap(2)
     print(map01)
 
     # CHECK: (d0, d1, d2, d3, d4) -> (d3, d4)
     map34 = map5.get_minor_submap(2)
     print(map34)
 
 
-run(testAffineMapDerive)
-
-
 # CHECK-LABEL: TEST: testAffineMapProperties
+@run
 def testAffineMapProperties():
   with Context():
     d0 = AffineDimExpr.get(0)
     d1 = AffineDimExpr.get(1)
     d2 = AffineDimExpr.get(2)
     map1 = AffineMap.get(3, 0, [d2, d0])
     map2 = AffineMap.get(3, 0, [d2, d0, d1])
     map3 = AffineMap.get(3, 1, [d2, d0, d1])
     # CHECK: False
     print(map1.is_permutation)
     # CHECK: True
     print(map1.is_projected_permutation)
     # CHECK: True
     print(map2.is_permutation)
     # CHECK: True
     print(map2.is_projected_permutation)
     # CHECK: False
     print(map3.is_permutation)
     # CHECK: False
     print(map3.is_projected_permutation)
 
 
-run(testAffineMapProperties)
-
-
 # CHECK-LABEL: TEST: testAffineMapExprs
+@run
 def testAffineMapExprs():
   with Context():
     d0 = AffineDimExpr.get(0)
     d1 = AffineDimExpr.get(1)
     d2 = AffineDimExpr.get(2)
     map3 = AffineMap.get(3, 1, [d2, d0, d1])
 
     # CHECK: 3
with Context():
     for expr in map3.results[-1:-4:-1]:
       # CHECK: d1
       # CHECK: d0
       # CHECK: d2
       print(expr)
     assert list(map3.results) == [d2, d0, d1]
 
 
-run(testAffineMapExprs)
-
-
 # CHECK-LABEL: TEST: testCompressUnusedSymbols
+@run
 def testCompressUnusedSymbols():
   with Context() as ctx:
     d0, d1, d2 = (AffineDimExpr.get(0), AffineDimExpr.get(1),
                   AffineDimExpr.get(2))
     s0, s1, s2 = (AffineSymbolExpr.get(0), AffineSymbolExpr.get(1),
                   AffineSymbolExpr.get(2))
     maps = [
         AffineMap.get(3, 3, [d2, d0, d1]),
with Context() as ctx:
     print(maps)
 
     #      CHECK: AffineMap((d0, d1, d2)[s0] -> (d2, d0, d1))
     # CHECK-SAME: AffineMap((d0, d1, d2)[s0] -> (d2, d0 + s0, d1))
     # CHECK-SAME: AffineMap((d0, d1, d2)[s0] -> (d1, d2, d0))
     print(compressed_maps)
 
 
-run(testCompressUnusedSymbols)
-
-
 # CHECK-LABEL: TEST: testReplace
+@run
 def testReplace():
   with Context() as ctx:
     d0, d1, d2 = (AffineDimExpr.get(0), AffineDimExpr.get(1),
                   AffineDimExpr.get(2))
     s0, s1, s2 = (AffineSymbolExpr.get(0), AffineSymbolExpr.get(1),
                   AffineSymbolExpr.get(2))
     map1 = AffineMap.get(3, 3, [d2, d0 + s1 + s2, d1 + s0])
 
     replace0 = map1.replace(s0, AffineConstantExpr.get(42), 3, 3)
     replace1 = map1.replace(s1, AffineConstantExpr.get(42), 3, 3)
     replace3 = map1.replace(s2, AffineConstantExpr.get(42), 3, 2)
 
     # CHECK: (d0, d1, d2)[s0, s1, s2] -> (d2, d0 + s1 + s2, d1 + 42)
     print(replace0)
 
     # CHECK: (d0, d1, d2)[s0, s1, s2] -> (d2, d0 + s2 + 42, d1 + s0)
     print(replace1)
 
     # CHECK: (d0, d1, d2)[s0, s1] -> (d2, d0 + s1 + 42, d1 + s0)
     print(replace3)
 
 
-run(testReplace)
+# CHECK-LABEL: TEST: testHash
+@run
+def testHash():
+  with Context():
+    d0, d1 = AffineDimExpr.get(0), AffineDimExpr.get(1)
+    m1 = AffineMap.get(2, 0, [d0, d1])
+    m2 = AffineMap.get(2, 0, [d1, d0])
+    assert hash(m1) == hash(AffineMap.get(2, 0, [d0, d1]))
+
+    dictionary = dict()
+    dictionary[m1] = 1
+    dictionary[m2] = 2
+    assert m1 in dictionary

mlir/test/python/ir/integer_set.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
+  return f
 
 
 # CHECK-LABEL: TEST: testIntegerSetCapsule
+@run
 def testIntegerSetCapsule():
   with Context() as ctx:
     is1 = IntegerSet.get_empty(1, 1, ctx)
   capsule = is1._CAPIPtr
   # CHECK: mlir.ir.IntegerSet._CAPIPtr
   print(capsule)
   is2 = IntegerSet._CAPICreate(capsule)
   assert is1 == is2
   assert is2.context is ctx
 
-run(testIntegerSetCapsule)
-
 
 # CHECK-LABEL: TEST: testIntegerSetGet
+@run
 def testIntegerSetGet():
   with Context():
     d0 = AffineDimExpr.get(0)
     d1 = AffineDimExpr.get(1)
     s0 = AffineSymbolExpr.get(0)
     c42 = AffineConstantExpr.get(42)
 
     # CHECK: (d0, d1)[s0] : (d0 - d1 == 0, s0 - 42 >= 0)
except RuntimeError as e:
       print(e)
 
     try:
       set0.get_replaced([d0, d1], [None], 1, 1)
     except RuntimeError as e:
       # CHECK: Invalid expression (None?) when attempting to create an IntegerSet by replacing symbols
       print(e)
 
-run(testIntegerSetGet)
-
 
 # CHECK-LABEL: TEST: testIntegerSetProperties
+@run
 def testIntegerSetProperties():
   with Context():
     d0 = AffineDimExpr.get(0)
     d1 = AffineDimExpr.get(1)
     s0 = AffineSymbolExpr.get(0)
     c42 = AffineConstantExpr.get(42)
 
     set0 = IntegerSet.get(2, 1, [d0 - d1, s0 - c42, s0 - d0], [True, False, False])
with Context():
 
     # CHECK-DAG: d0 - d1 == 0
     # CHECK-DAG: s0 - 42 >= 0
     # CHECK-DAG: -d0 + s0 >= 0
     for cstr in set0.constraints:
       print(cstr.expr, end='')
       print(" == 0" if cstr.is_eq else " >= 0")
 
-run(testIntegerSetProperties)
+
+# CHECK_LABEL: TEST: testHash
+@run
+def testHash():
+  with Context():
+    d0 = AffineDimExpr.get(0)
+    d1 = AffineDimExpr.get(1)
+    set = IntegerSet.get(2, 0, [d0 + d1], [True])
+
+    assert hash(set) == hash(IntegerSet.get(2, 0, [d0 + d1], [True]))
+
+    dictionary = dict()
+    dictionary[set] = 42
+    assert set in dictionary