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);
}
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);
}
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);
}
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);
}
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);
}
.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) {
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__",
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.")
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;
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(); })
# 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
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
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
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
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
# 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():
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
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)
assert am2.context is ctx
-run(testAffineMapCapsule)
-
-
# CHECK-LABEL: TEST: testAffineMapGet
+@run
def testAffineMapGet():
with Context() as ctx:
d0 = AffineDimExpr.get(0)
print(e)
-run(testAffineMapGet)
-
-
# CHECK-LABEL: TEST: testAffineMapDerive
+@run
def testAffineMapDerive():
with Context() as ctx:
map5 = AffineMap.get_identity(5)
print(map34)
-run(testAffineMapDerive)
-
-
# CHECK-LABEL: TEST: testAffineMapProperties
+@run
def testAffineMapProperties():
with Context():
d0 = AffineDimExpr.get(0)
print(map3.is_projected_permutation)
-run(testAffineMapProperties)
-
-
# CHECK-LABEL: TEST: testAffineMapExprs
+@run
def testAffineMapExprs():
with Context():
d0 = AffineDimExpr.get(0)
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),
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),
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