1 //===- IRAffine.cpp - Exports 'ir' module affine related bindings ---------===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 9 #include <utility> 10 11 #include "IRModule.h" 12 13 #include "PybindUtils.h" 14 15 #include "mlir-c/AffineMap.h" 16 #include "mlir-c/Bindings/Python/Interop.h" 17 #include "mlir-c/IntegerSet.h" 18 19 namespace py = pybind11; 20 using namespace mlir; 21 using namespace mlir::python; 22 23 using llvm::SmallVector; 24 using llvm::StringRef; 25 using llvm::Twine; 26 27 static const char kDumpDocstring[] = 28 R"(Dumps a debug representation of the object to stderr.)"; 29 30 /// Attempts to populate `result` with the content of `list` casted to the 31 /// appropriate type (Python and C types are provided as template arguments). 32 /// Throws errors in case of failure, using "action" to describe what the caller 33 /// was attempting to do. 34 template <typename PyType, typename CType> 35 static void pyListToVector(const py::list &list, 36 llvm::SmallVectorImpl<CType> &result, 37 StringRef action) { 38 result.reserve(py::len(list)); 39 for (py::handle item : list) { 40 try { 41 result.push_back(item.cast<PyType>()); 42 } catch (py::cast_error &err) { 43 std::string msg = (llvm::Twine("Invalid expression when ") + action + 44 " (" + err.what() + ")") 45 .str(); 46 throw py::cast_error(msg); 47 } catch (py::reference_cast_error &err) { 48 std::string msg = (llvm::Twine("Invalid expression (None?) when ") + 49 action + " (" + err.what() + ")") 50 .str(); 51 throw py::cast_error(msg); 52 } 53 } 54 } 55 56 template <typename PermutationTy> 57 static bool isPermutation(std::vector<PermutationTy> permutation) { 58 llvm::SmallVector<bool, 8> seen(permutation.size(), false); 59 for (auto val : permutation) { 60 if (val < permutation.size()) { 61 if (seen[val]) 62 return false; 63 seen[val] = true; 64 continue; 65 } 66 return false; 67 } 68 return true; 69 } 70 71 namespace { 72 73 /// CRTP base class for Python MLIR affine expressions that subclass AffineExpr 74 /// and should be castable from it. Intermediate hierarchy classes can be 75 /// modeled by specifying BaseTy. 76 template <typename DerivedTy, typename BaseTy = PyAffineExpr> 77 class PyConcreteAffineExpr : public BaseTy { 78 public: 79 // Derived classes must define statics for: 80 // IsAFunctionTy isaFunction 81 // const char *pyClassName 82 // and redefine bindDerived. 83 using ClassTy = py::class_<DerivedTy, BaseTy>; 84 using IsAFunctionTy = bool (*)(MlirAffineExpr); 85 86 PyConcreteAffineExpr() = default; 87 PyConcreteAffineExpr(PyMlirContextRef contextRef, MlirAffineExpr affineExpr) 88 : BaseTy(std::move(contextRef), affineExpr) {} 89 PyConcreteAffineExpr(PyAffineExpr &orig) 90 : PyConcreteAffineExpr(orig.getContext(), castFrom(orig)) {} 91 92 static MlirAffineExpr castFrom(PyAffineExpr &orig) { 93 if (!DerivedTy::isaFunction(orig)) { 94 auto origRepr = py::repr(py::cast(orig)).cast<std::string>(); 95 throw SetPyError(PyExc_ValueError, 96 Twine("Cannot cast affine expression to ") + 97 DerivedTy::pyClassName + " (from " + origRepr + ")"); 98 } 99 return orig; 100 } 101 102 static void bind(py::module &m) { 103 auto cls = ClassTy(m, DerivedTy::pyClassName, py::module_local()); 104 cls.def(py::init<PyAffineExpr &>(), py::arg("expr")); 105 cls.def_static( 106 "isinstance", 107 [](PyAffineExpr &otherAffineExpr) -> bool { 108 return DerivedTy::isaFunction(otherAffineExpr); 109 }, 110 py::arg("other")); 111 DerivedTy::bindDerived(cls); 112 } 113 114 /// Implemented by derived classes to add methods to the Python subclass. 115 static void bindDerived(ClassTy &m) {} 116 }; 117 118 class PyAffineConstantExpr : public PyConcreteAffineExpr<PyAffineConstantExpr> { 119 public: 120 static constexpr IsAFunctionTy isaFunction = mlirAffineExprIsAConstant; 121 static constexpr const char *pyClassName = "AffineConstantExpr"; 122 using PyConcreteAffineExpr::PyConcreteAffineExpr; 123 124 static PyAffineConstantExpr get(intptr_t value, 125 DefaultingPyMlirContext context) { 126 MlirAffineExpr affineExpr = 127 mlirAffineConstantExprGet(context->get(), static_cast<int64_t>(value)); 128 return PyAffineConstantExpr(context->getRef(), affineExpr); 129 } 130 131 static void bindDerived(ClassTy &c) { 132 c.def_static("get", &PyAffineConstantExpr::get, py::arg("value"), 133 py::arg("context") = py::none()); 134 c.def_property_readonly("value", [](PyAffineConstantExpr &self) { 135 return mlirAffineConstantExprGetValue(self); 136 }); 137 } 138 }; 139 140 class PyAffineDimExpr : public PyConcreteAffineExpr<PyAffineDimExpr> { 141 public: 142 static constexpr IsAFunctionTy isaFunction = mlirAffineExprIsADim; 143 static constexpr const char *pyClassName = "AffineDimExpr"; 144 using PyConcreteAffineExpr::PyConcreteAffineExpr; 145 146 static PyAffineDimExpr get(intptr_t pos, DefaultingPyMlirContext context) { 147 MlirAffineExpr affineExpr = mlirAffineDimExprGet(context->get(), pos); 148 return PyAffineDimExpr(context->getRef(), affineExpr); 149 } 150 151 static void bindDerived(ClassTy &c) { 152 c.def_static("get", &PyAffineDimExpr::get, py::arg("position"), 153 py::arg("context") = py::none()); 154 c.def_property_readonly("position", [](PyAffineDimExpr &self) { 155 return mlirAffineDimExprGetPosition(self); 156 }); 157 } 158 }; 159 160 class PyAffineSymbolExpr : public PyConcreteAffineExpr<PyAffineSymbolExpr> { 161 public: 162 static constexpr IsAFunctionTy isaFunction = mlirAffineExprIsASymbol; 163 static constexpr const char *pyClassName = "AffineSymbolExpr"; 164 using PyConcreteAffineExpr::PyConcreteAffineExpr; 165 166 static PyAffineSymbolExpr get(intptr_t pos, DefaultingPyMlirContext context) { 167 MlirAffineExpr affineExpr = mlirAffineSymbolExprGet(context->get(), pos); 168 return PyAffineSymbolExpr(context->getRef(), affineExpr); 169 } 170 171 static void bindDerived(ClassTy &c) { 172 c.def_static("get", &PyAffineSymbolExpr::get, py::arg("position"), 173 py::arg("context") = py::none()); 174 c.def_property_readonly("position", [](PyAffineSymbolExpr &self) { 175 return mlirAffineSymbolExprGetPosition(self); 176 }); 177 } 178 }; 179 180 class PyAffineBinaryExpr : public PyConcreteAffineExpr<PyAffineBinaryExpr> { 181 public: 182 static constexpr IsAFunctionTy isaFunction = mlirAffineExprIsABinary; 183 static constexpr const char *pyClassName = "AffineBinaryExpr"; 184 using PyConcreteAffineExpr::PyConcreteAffineExpr; 185 186 PyAffineExpr lhs() { 187 MlirAffineExpr lhsExpr = mlirAffineBinaryOpExprGetLHS(get()); 188 return PyAffineExpr(getContext(), lhsExpr); 189 } 190 191 PyAffineExpr rhs() { 192 MlirAffineExpr rhsExpr = mlirAffineBinaryOpExprGetRHS(get()); 193 return PyAffineExpr(getContext(), rhsExpr); 194 } 195 196 static void bindDerived(ClassTy &c) { 197 c.def_property_readonly("lhs", &PyAffineBinaryExpr::lhs); 198 c.def_property_readonly("rhs", &PyAffineBinaryExpr::rhs); 199 } 200 }; 201 202 class PyAffineAddExpr 203 : public PyConcreteAffineExpr<PyAffineAddExpr, PyAffineBinaryExpr> { 204 public: 205 static constexpr IsAFunctionTy isaFunction = mlirAffineExprIsAAdd; 206 static constexpr const char *pyClassName = "AffineAddExpr"; 207 using PyConcreteAffineExpr::PyConcreteAffineExpr; 208 209 static PyAffineAddExpr get(PyAffineExpr lhs, const PyAffineExpr &rhs) { 210 MlirAffineExpr expr = mlirAffineAddExprGet(lhs, rhs); 211 return PyAffineAddExpr(lhs.getContext(), expr); 212 } 213 214 static PyAffineAddExpr getRHSConstant(PyAffineExpr lhs, intptr_t rhs) { 215 MlirAffineExpr expr = mlirAffineAddExprGet( 216 lhs, mlirAffineConstantExprGet(mlirAffineExprGetContext(lhs), rhs)); 217 return PyAffineAddExpr(lhs.getContext(), expr); 218 } 219 220 static PyAffineAddExpr getLHSConstant(intptr_t lhs, PyAffineExpr rhs) { 221 MlirAffineExpr expr = mlirAffineAddExprGet( 222 mlirAffineConstantExprGet(mlirAffineExprGetContext(rhs), lhs), rhs); 223 return PyAffineAddExpr(rhs.getContext(), expr); 224 } 225 226 static void bindDerived(ClassTy &c) { 227 c.def_static("get", &PyAffineAddExpr::get); 228 } 229 }; 230 231 class PyAffineMulExpr 232 : public PyConcreteAffineExpr<PyAffineMulExpr, PyAffineBinaryExpr> { 233 public: 234 static constexpr IsAFunctionTy isaFunction = mlirAffineExprIsAMul; 235 static constexpr const char *pyClassName = "AffineMulExpr"; 236 using PyConcreteAffineExpr::PyConcreteAffineExpr; 237 238 static PyAffineMulExpr get(PyAffineExpr lhs, const PyAffineExpr &rhs) { 239 MlirAffineExpr expr = mlirAffineMulExprGet(lhs, rhs); 240 return PyAffineMulExpr(lhs.getContext(), expr); 241 } 242 243 static PyAffineMulExpr getRHSConstant(PyAffineExpr lhs, intptr_t rhs) { 244 MlirAffineExpr expr = mlirAffineMulExprGet( 245 lhs, mlirAffineConstantExprGet(mlirAffineExprGetContext(lhs), rhs)); 246 return PyAffineMulExpr(lhs.getContext(), expr); 247 } 248 249 static PyAffineMulExpr getLHSConstant(intptr_t lhs, PyAffineExpr rhs) { 250 MlirAffineExpr expr = mlirAffineMulExprGet( 251 mlirAffineConstantExprGet(mlirAffineExprGetContext(rhs), lhs), rhs); 252 return PyAffineMulExpr(rhs.getContext(), expr); 253 } 254 255 static void bindDerived(ClassTy &c) { 256 c.def_static("get", &PyAffineMulExpr::get); 257 } 258 }; 259 260 class PyAffineModExpr 261 : public PyConcreteAffineExpr<PyAffineModExpr, PyAffineBinaryExpr> { 262 public: 263 static constexpr IsAFunctionTy isaFunction = mlirAffineExprIsAMod; 264 static constexpr const char *pyClassName = "AffineModExpr"; 265 using PyConcreteAffineExpr::PyConcreteAffineExpr; 266 267 static PyAffineModExpr get(PyAffineExpr lhs, const PyAffineExpr &rhs) { 268 MlirAffineExpr expr = mlirAffineModExprGet(lhs, rhs); 269 return PyAffineModExpr(lhs.getContext(), expr); 270 } 271 272 static PyAffineModExpr getRHSConstant(PyAffineExpr lhs, intptr_t rhs) { 273 MlirAffineExpr expr = mlirAffineModExprGet( 274 lhs, mlirAffineConstantExprGet(mlirAffineExprGetContext(lhs), rhs)); 275 return PyAffineModExpr(lhs.getContext(), expr); 276 } 277 278 static PyAffineModExpr getLHSConstant(intptr_t lhs, PyAffineExpr rhs) { 279 MlirAffineExpr expr = mlirAffineModExprGet( 280 mlirAffineConstantExprGet(mlirAffineExprGetContext(rhs), lhs), rhs); 281 return PyAffineModExpr(rhs.getContext(), expr); 282 } 283 284 static void bindDerived(ClassTy &c) { 285 c.def_static("get", &PyAffineModExpr::get); 286 } 287 }; 288 289 class PyAffineFloorDivExpr 290 : public PyConcreteAffineExpr<PyAffineFloorDivExpr, PyAffineBinaryExpr> { 291 public: 292 static constexpr IsAFunctionTy isaFunction = mlirAffineExprIsAFloorDiv; 293 static constexpr const char *pyClassName = "AffineFloorDivExpr"; 294 using PyConcreteAffineExpr::PyConcreteAffineExpr; 295 296 static PyAffineFloorDivExpr get(PyAffineExpr lhs, const PyAffineExpr &rhs) { 297 MlirAffineExpr expr = mlirAffineFloorDivExprGet(lhs, rhs); 298 return PyAffineFloorDivExpr(lhs.getContext(), expr); 299 } 300 301 static PyAffineFloorDivExpr getRHSConstant(PyAffineExpr lhs, intptr_t rhs) { 302 MlirAffineExpr expr = mlirAffineFloorDivExprGet( 303 lhs, mlirAffineConstantExprGet(mlirAffineExprGetContext(lhs), rhs)); 304 return PyAffineFloorDivExpr(lhs.getContext(), expr); 305 } 306 307 static PyAffineFloorDivExpr getLHSConstant(intptr_t lhs, PyAffineExpr rhs) { 308 MlirAffineExpr expr = mlirAffineFloorDivExprGet( 309 mlirAffineConstantExprGet(mlirAffineExprGetContext(rhs), lhs), rhs); 310 return PyAffineFloorDivExpr(rhs.getContext(), expr); 311 } 312 313 static void bindDerived(ClassTy &c) { 314 c.def_static("get", &PyAffineFloorDivExpr::get); 315 } 316 }; 317 318 class PyAffineCeilDivExpr 319 : public PyConcreteAffineExpr<PyAffineCeilDivExpr, PyAffineBinaryExpr> { 320 public: 321 static constexpr IsAFunctionTy isaFunction = mlirAffineExprIsACeilDiv; 322 static constexpr const char *pyClassName = "AffineCeilDivExpr"; 323 using PyConcreteAffineExpr::PyConcreteAffineExpr; 324 325 static PyAffineCeilDivExpr get(PyAffineExpr lhs, const PyAffineExpr &rhs) { 326 MlirAffineExpr expr = mlirAffineCeilDivExprGet(lhs, rhs); 327 return PyAffineCeilDivExpr(lhs.getContext(), expr); 328 } 329 330 static PyAffineCeilDivExpr getRHSConstant(PyAffineExpr lhs, intptr_t rhs) { 331 MlirAffineExpr expr = mlirAffineCeilDivExprGet( 332 lhs, mlirAffineConstantExprGet(mlirAffineExprGetContext(lhs), rhs)); 333 return PyAffineCeilDivExpr(lhs.getContext(), expr); 334 } 335 336 static PyAffineCeilDivExpr getLHSConstant(intptr_t lhs, PyAffineExpr rhs) { 337 MlirAffineExpr expr = mlirAffineCeilDivExprGet( 338 mlirAffineConstantExprGet(mlirAffineExprGetContext(rhs), lhs), rhs); 339 return PyAffineCeilDivExpr(rhs.getContext(), expr); 340 } 341 342 static void bindDerived(ClassTy &c) { 343 c.def_static("get", &PyAffineCeilDivExpr::get); 344 } 345 }; 346 347 } // namespace 348 349 bool PyAffineExpr::operator==(const PyAffineExpr &other) { 350 return mlirAffineExprEqual(affineExpr, other.affineExpr); 351 } 352 353 py::object PyAffineExpr::getCapsule() { 354 return py::reinterpret_steal<py::object>( 355 mlirPythonAffineExprToCapsule(*this)); 356 } 357 358 PyAffineExpr PyAffineExpr::createFromCapsule(py::object capsule) { 359 MlirAffineExpr rawAffineExpr = mlirPythonCapsuleToAffineExpr(capsule.ptr()); 360 if (mlirAffineExprIsNull(rawAffineExpr)) 361 throw py::error_already_set(); 362 return PyAffineExpr( 363 PyMlirContext::forContext(mlirAffineExprGetContext(rawAffineExpr)), 364 rawAffineExpr); 365 } 366 367 //------------------------------------------------------------------------------ 368 // PyAffineMap and utilities. 369 //------------------------------------------------------------------------------ 370 namespace { 371 372 /// A list of expressions contained in an affine map. Internally these are 373 /// stored as a consecutive array leading to inexpensive random access. Both 374 /// the map and the expression are owned by the context so we need not bother 375 /// with lifetime extension. 376 class PyAffineMapExprList 377 : public Sliceable<PyAffineMapExprList, PyAffineExpr> { 378 public: 379 static constexpr const char *pyClassName = "AffineExprList"; 380 381 PyAffineMapExprList(const PyAffineMap &map, intptr_t startIndex = 0, 382 intptr_t length = -1, intptr_t step = 1) 383 : Sliceable(startIndex, 384 length == -1 ? mlirAffineMapGetNumResults(map) : length, 385 step), 386 affineMap(map) {} 387 388 intptr_t getNumElements() { return mlirAffineMapGetNumResults(affineMap); } 389 390 PyAffineExpr getElement(intptr_t pos) { 391 return PyAffineExpr(affineMap.getContext(), 392 mlirAffineMapGetResult(affineMap, pos)); 393 } 394 395 PyAffineMapExprList slice(intptr_t startIndex, intptr_t length, 396 intptr_t step) { 397 return PyAffineMapExprList(affineMap, startIndex, length, step); 398 } 399 400 private: 401 PyAffineMap affineMap; 402 }; 403 } // namespace 404 405 bool PyAffineMap::operator==(const PyAffineMap &other) { 406 return mlirAffineMapEqual(affineMap, other.affineMap); 407 } 408 409 py::object PyAffineMap::getCapsule() { 410 return py::reinterpret_steal<py::object>(mlirPythonAffineMapToCapsule(*this)); 411 } 412 413 PyAffineMap PyAffineMap::createFromCapsule(py::object capsule) { 414 MlirAffineMap rawAffineMap = mlirPythonCapsuleToAffineMap(capsule.ptr()); 415 if (mlirAffineMapIsNull(rawAffineMap)) 416 throw py::error_already_set(); 417 return PyAffineMap( 418 PyMlirContext::forContext(mlirAffineMapGetContext(rawAffineMap)), 419 rawAffineMap); 420 } 421 422 //------------------------------------------------------------------------------ 423 // PyIntegerSet and utilities. 424 //------------------------------------------------------------------------------ 425 namespace { 426 427 class PyIntegerSetConstraint { 428 public: 429 PyIntegerSetConstraint(PyIntegerSet set, intptr_t pos) 430 : set(std::move(set)), pos(pos) {} 431 432 PyAffineExpr getExpr() { 433 return PyAffineExpr(set.getContext(), 434 mlirIntegerSetGetConstraint(set, pos)); 435 } 436 437 bool isEq() { return mlirIntegerSetIsConstraintEq(set, pos); } 438 439 static void bind(py::module &m) { 440 py::class_<PyIntegerSetConstraint>(m, "IntegerSetConstraint", 441 py::module_local()) 442 .def_property_readonly("expr", &PyIntegerSetConstraint::getExpr) 443 .def_property_readonly("is_eq", &PyIntegerSetConstraint::isEq); 444 } 445 446 private: 447 PyIntegerSet set; 448 intptr_t pos; 449 }; 450 451 class PyIntegerSetConstraintList 452 : public Sliceable<PyIntegerSetConstraintList, PyIntegerSetConstraint> { 453 public: 454 static constexpr const char *pyClassName = "IntegerSetConstraintList"; 455 456 PyIntegerSetConstraintList(const PyIntegerSet &set, intptr_t startIndex = 0, 457 intptr_t length = -1, intptr_t step = 1) 458 : Sliceable(startIndex, 459 length == -1 ? mlirIntegerSetGetNumConstraints(set) : length, 460 step), 461 set(set) {} 462 463 intptr_t getNumElements() { return mlirIntegerSetGetNumConstraints(set); } 464 465 PyIntegerSetConstraint getElement(intptr_t pos) { 466 return PyIntegerSetConstraint(set, pos); 467 } 468 469 PyIntegerSetConstraintList slice(intptr_t startIndex, intptr_t length, 470 intptr_t step) { 471 return PyIntegerSetConstraintList(set, startIndex, length, step); 472 } 473 474 private: 475 PyIntegerSet set; 476 }; 477 } // namespace 478 479 bool PyIntegerSet::operator==(const PyIntegerSet &other) { 480 return mlirIntegerSetEqual(integerSet, other.integerSet); 481 } 482 483 py::object PyIntegerSet::getCapsule() { 484 return py::reinterpret_steal<py::object>( 485 mlirPythonIntegerSetToCapsule(*this)); 486 } 487 488 PyIntegerSet PyIntegerSet::createFromCapsule(py::object capsule) { 489 MlirIntegerSet rawIntegerSet = mlirPythonCapsuleToIntegerSet(capsule.ptr()); 490 if (mlirIntegerSetIsNull(rawIntegerSet)) 491 throw py::error_already_set(); 492 return PyIntegerSet( 493 PyMlirContext::forContext(mlirIntegerSetGetContext(rawIntegerSet)), 494 rawIntegerSet); 495 } 496 497 void mlir::python::populateIRAffine(py::module &m) { 498 //---------------------------------------------------------------------------- 499 // Mapping of PyAffineExpr and derived classes. 500 //---------------------------------------------------------------------------- 501 py::class_<PyAffineExpr>(m, "AffineExpr", py::module_local()) 502 .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, 503 &PyAffineExpr::getCapsule) 504 .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyAffineExpr::createFromCapsule) 505 .def("__add__", &PyAffineAddExpr::get) 506 .def("__add__", &PyAffineAddExpr::getRHSConstant) 507 .def("__radd__", &PyAffineAddExpr::getRHSConstant) 508 .def("__mul__", &PyAffineMulExpr::get) 509 .def("__mul__", &PyAffineMulExpr::getRHSConstant) 510 .def("__rmul__", &PyAffineMulExpr::getRHSConstant) 511 .def("__mod__", &PyAffineModExpr::get) 512 .def("__mod__", &PyAffineModExpr::getRHSConstant) 513 .def("__rmod__", 514 [](PyAffineExpr &self, intptr_t other) { 515 return PyAffineModExpr::get( 516 PyAffineConstantExpr::get(other, *self.getContext().get()), 517 self); 518 }) 519 .def("__sub__", 520 [](PyAffineExpr &self, PyAffineExpr &other) { 521 auto negOne = 522 PyAffineConstantExpr::get(-1, *self.getContext().get()); 523 return PyAffineAddExpr::get(self, 524 PyAffineMulExpr::get(negOne, other)); 525 }) 526 .def("__sub__", 527 [](PyAffineExpr &self, intptr_t other) { 528 return PyAffineAddExpr::get( 529 self, 530 PyAffineConstantExpr::get(-other, *self.getContext().get())); 531 }) 532 .def("__rsub__", 533 [](PyAffineExpr &self, intptr_t other) { 534 return PyAffineAddExpr::getLHSConstant( 535 other, PyAffineMulExpr::getLHSConstant(-1, self)); 536 }) 537 .def("__eq__", [](PyAffineExpr &self, 538 PyAffineExpr &other) { return self == other; }) 539 .def("__eq__", 540 [](PyAffineExpr &self, py::object &other) { return false; }) 541 .def("__str__", 542 [](PyAffineExpr &self) { 543 PyPrintAccumulator printAccum; 544 mlirAffineExprPrint(self, printAccum.getCallback(), 545 printAccum.getUserData()); 546 return printAccum.join(); 547 }) 548 .def("__repr__", 549 [](PyAffineExpr &self) { 550 PyPrintAccumulator printAccum; 551 printAccum.parts.append("AffineExpr("); 552 mlirAffineExprPrint(self, printAccum.getCallback(), 553 printAccum.getUserData()); 554 printAccum.parts.append(")"); 555 return printAccum.join(); 556 }) 557 .def("__hash__", 558 [](PyAffineExpr &self) { 559 return static_cast<size_t>(llvm::hash_value(self.get().ptr)); 560 }) 561 .def_property_readonly( 562 "context", 563 [](PyAffineExpr &self) { return self.getContext().getObject(); }) 564 .def("compose", 565 [](PyAffineExpr &self, PyAffineMap &other) { 566 return PyAffineExpr(self.getContext(), 567 mlirAffineExprCompose(self, other)); 568 }) 569 .def_static( 570 "get_add", &PyAffineAddExpr::get, 571 "Gets an affine expression containing a sum of two expressions.") 572 .def_static("get_add", &PyAffineAddExpr::getLHSConstant, 573 "Gets an affine expression containing a sum of a constant " 574 "and another expression.") 575 .def_static("get_add", &PyAffineAddExpr::getRHSConstant, 576 "Gets an affine expression containing a sum of an expression " 577 "and a constant.") 578 .def_static( 579 "get_mul", &PyAffineMulExpr::get, 580 "Gets an affine expression containing a product of two expressions.") 581 .def_static("get_mul", &PyAffineMulExpr::getLHSConstant, 582 "Gets an affine expression containing a product of a " 583 "constant and another expression.") 584 .def_static("get_mul", &PyAffineMulExpr::getRHSConstant, 585 "Gets an affine expression containing a product of an " 586 "expression and a constant.") 587 .def_static("get_mod", &PyAffineModExpr::get, 588 "Gets an affine expression containing the modulo of dividing " 589 "one expression by another.") 590 .def_static("get_mod", &PyAffineModExpr::getLHSConstant, 591 "Gets a semi-affine expression containing the modulo of " 592 "dividing a constant by an expression.") 593 .def_static("get_mod", &PyAffineModExpr::getRHSConstant, 594 "Gets an affine expression containing the module of dividing" 595 "an expression by a constant.") 596 .def_static("get_floor_div", &PyAffineFloorDivExpr::get, 597 "Gets an affine expression containing the rounded-down " 598 "result of dividing one expression by another.") 599 .def_static("get_floor_div", &PyAffineFloorDivExpr::getLHSConstant, 600 "Gets a semi-affine expression containing the rounded-down " 601 "result of dividing a constant by an expression.") 602 .def_static("get_floor_div", &PyAffineFloorDivExpr::getRHSConstant, 603 "Gets an affine expression containing the rounded-down " 604 "result of dividing an expression by a constant.") 605 .def_static("get_ceil_div", &PyAffineCeilDivExpr::get, 606 "Gets an affine expression containing the rounded-up result " 607 "of dividing one expression by another.") 608 .def_static("get_ceil_div", &PyAffineCeilDivExpr::getLHSConstant, 609 "Gets a semi-affine expression containing the rounded-up " 610 "result of dividing a constant by an expression.") 611 .def_static("get_ceil_div", &PyAffineCeilDivExpr::getRHSConstant, 612 "Gets an affine expression containing the rounded-up result " 613 "of dividing an expression by a constant.") 614 .def_static("get_constant", &PyAffineConstantExpr::get, py::arg("value"), 615 py::arg("context") = py::none(), 616 "Gets a constant affine expression with the given value.") 617 .def_static( 618 "get_dim", &PyAffineDimExpr::get, py::arg("position"), 619 py::arg("context") = py::none(), 620 "Gets an affine expression of a dimension at the given position.") 621 .def_static( 622 "get_symbol", &PyAffineSymbolExpr::get, py::arg("position"), 623 py::arg("context") = py::none(), 624 "Gets an affine expression of a symbol at the given position.") 625 .def( 626 "dump", [](PyAffineExpr &self) { mlirAffineExprDump(self); }, 627 kDumpDocstring); 628 PyAffineConstantExpr::bind(m); 629 PyAffineDimExpr::bind(m); 630 PyAffineSymbolExpr::bind(m); 631 PyAffineBinaryExpr::bind(m); 632 PyAffineAddExpr::bind(m); 633 PyAffineMulExpr::bind(m); 634 PyAffineModExpr::bind(m); 635 PyAffineFloorDivExpr::bind(m); 636 PyAffineCeilDivExpr::bind(m); 637 638 //---------------------------------------------------------------------------- 639 // Mapping of PyAffineMap. 640 //---------------------------------------------------------------------------- 641 py::class_<PyAffineMap>(m, "AffineMap", py::module_local()) 642 .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, 643 &PyAffineMap::getCapsule) 644 .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyAffineMap::createFromCapsule) 645 .def("__eq__", 646 [](PyAffineMap &self, PyAffineMap &other) { return self == other; }) 647 .def("__eq__", [](PyAffineMap &self, py::object &other) { return false; }) 648 .def("__str__", 649 [](PyAffineMap &self) { 650 PyPrintAccumulator printAccum; 651 mlirAffineMapPrint(self, printAccum.getCallback(), 652 printAccum.getUserData()); 653 return printAccum.join(); 654 }) 655 .def("__repr__", 656 [](PyAffineMap &self) { 657 PyPrintAccumulator printAccum; 658 printAccum.parts.append("AffineMap("); 659 mlirAffineMapPrint(self, printAccum.getCallback(), 660 printAccum.getUserData()); 661 printAccum.parts.append(")"); 662 return printAccum.join(); 663 }) 664 .def("__hash__", 665 [](PyAffineMap &self) { 666 return static_cast<size_t>(llvm::hash_value(self.get().ptr)); 667 }) 668 .def_static("compress_unused_symbols", 669 [](py::list affineMaps, DefaultingPyMlirContext context) { 670 SmallVector<MlirAffineMap> maps; 671 pyListToVector<PyAffineMap, MlirAffineMap>( 672 affineMaps, maps, "attempting to create an AffineMap"); 673 std::vector<MlirAffineMap> compressed(affineMaps.size()); 674 auto populate = [](void *result, intptr_t idx, 675 MlirAffineMap m) { 676 static_cast<MlirAffineMap *>(result)[idx] = (m); 677 }; 678 mlirAffineMapCompressUnusedSymbols( 679 maps.data(), maps.size(), compressed.data(), populate); 680 std::vector<PyAffineMap> res; 681 res.reserve(compressed.size()); 682 for (auto m : compressed) 683 res.emplace_back(context->getRef(), m); 684 return res; 685 }) 686 .def_property_readonly( 687 "context", 688 [](PyAffineMap &self) { return self.getContext().getObject(); }, 689 "Context that owns the Affine Map") 690 .def( 691 "dump", [](PyAffineMap &self) { mlirAffineMapDump(self); }, 692 kDumpDocstring) 693 .def_static( 694 "get", 695 [](intptr_t dimCount, intptr_t symbolCount, py::list exprs, 696 DefaultingPyMlirContext context) { 697 SmallVector<MlirAffineExpr> affineExprs; 698 pyListToVector<PyAffineExpr, MlirAffineExpr>( 699 exprs, affineExprs, "attempting to create an AffineMap"); 700 MlirAffineMap map = 701 mlirAffineMapGet(context->get(), dimCount, symbolCount, 702 affineExprs.size(), affineExprs.data()); 703 return PyAffineMap(context->getRef(), map); 704 }, 705 py::arg("dim_count"), py::arg("symbol_count"), py::arg("exprs"), 706 py::arg("context") = py::none(), 707 "Gets a map with the given expressions as results.") 708 .def_static( 709 "get_constant", 710 [](intptr_t value, DefaultingPyMlirContext context) { 711 MlirAffineMap affineMap = 712 mlirAffineMapConstantGet(context->get(), value); 713 return PyAffineMap(context->getRef(), affineMap); 714 }, 715 py::arg("value"), py::arg("context") = py::none(), 716 "Gets an affine map with a single constant result") 717 .def_static( 718 "get_empty", 719 [](DefaultingPyMlirContext context) { 720 MlirAffineMap affineMap = mlirAffineMapEmptyGet(context->get()); 721 return PyAffineMap(context->getRef(), affineMap); 722 }, 723 py::arg("context") = py::none(), "Gets an empty affine map.") 724 .def_static( 725 "get_identity", 726 [](intptr_t nDims, DefaultingPyMlirContext context) { 727 MlirAffineMap affineMap = 728 mlirAffineMapMultiDimIdentityGet(context->get(), nDims); 729 return PyAffineMap(context->getRef(), affineMap); 730 }, 731 py::arg("n_dims"), py::arg("context") = py::none(), 732 "Gets an identity map with the given number of dimensions.") 733 .def_static( 734 "get_minor_identity", 735 [](intptr_t nDims, intptr_t nResults, 736 DefaultingPyMlirContext context) { 737 MlirAffineMap affineMap = 738 mlirAffineMapMinorIdentityGet(context->get(), nDims, nResults); 739 return PyAffineMap(context->getRef(), affineMap); 740 }, 741 py::arg("n_dims"), py::arg("n_results"), 742 py::arg("context") = py::none(), 743 "Gets a minor identity map with the given number of dimensions and " 744 "results.") 745 .def_static( 746 "get_permutation", 747 [](std::vector<unsigned> permutation, 748 DefaultingPyMlirContext context) { 749 if (!isPermutation(permutation)) 750 throw py::cast_error("Invalid permutation when attempting to " 751 "create an AffineMap"); 752 MlirAffineMap affineMap = mlirAffineMapPermutationGet( 753 context->get(), permutation.size(), permutation.data()); 754 return PyAffineMap(context->getRef(), affineMap); 755 }, 756 py::arg("permutation"), py::arg("context") = py::none(), 757 "Gets an affine map that permutes its inputs.") 758 .def( 759 "get_submap", 760 [](PyAffineMap &self, std::vector<intptr_t> &resultPos) { 761 intptr_t numResults = mlirAffineMapGetNumResults(self); 762 for (intptr_t pos : resultPos) { 763 if (pos < 0 || pos >= numResults) 764 throw py::value_error("result position out of bounds"); 765 } 766 MlirAffineMap affineMap = mlirAffineMapGetSubMap( 767 self, resultPos.size(), resultPos.data()); 768 return PyAffineMap(self.getContext(), affineMap); 769 }, 770 py::arg("result_positions")) 771 .def( 772 "get_major_submap", 773 [](PyAffineMap &self, intptr_t nResults) { 774 if (nResults >= mlirAffineMapGetNumResults(self)) 775 throw py::value_error("number of results out of bounds"); 776 MlirAffineMap affineMap = 777 mlirAffineMapGetMajorSubMap(self, nResults); 778 return PyAffineMap(self.getContext(), affineMap); 779 }, 780 py::arg("n_results")) 781 .def( 782 "get_minor_submap", 783 [](PyAffineMap &self, intptr_t nResults) { 784 if (nResults >= mlirAffineMapGetNumResults(self)) 785 throw py::value_error("number of results out of bounds"); 786 MlirAffineMap affineMap = 787 mlirAffineMapGetMinorSubMap(self, nResults); 788 return PyAffineMap(self.getContext(), affineMap); 789 }, 790 py::arg("n_results")) 791 .def( 792 "replace", 793 [](PyAffineMap &self, PyAffineExpr &expression, 794 PyAffineExpr &replacement, intptr_t numResultDims, 795 intptr_t numResultSyms) { 796 MlirAffineMap affineMap = mlirAffineMapReplace( 797 self, expression, replacement, numResultDims, numResultSyms); 798 return PyAffineMap(self.getContext(), affineMap); 799 }, 800 py::arg("expr"), py::arg("replacement"), py::arg("n_result_dims"), 801 py::arg("n_result_syms")) 802 .def_property_readonly( 803 "is_permutation", 804 [](PyAffineMap &self) { return mlirAffineMapIsPermutation(self); }) 805 .def_property_readonly("is_projected_permutation", 806 [](PyAffineMap &self) { 807 return mlirAffineMapIsProjectedPermutation(self); 808 }) 809 .def_property_readonly( 810 "n_dims", 811 [](PyAffineMap &self) { return mlirAffineMapGetNumDims(self); }) 812 .def_property_readonly( 813 "n_inputs", 814 [](PyAffineMap &self) { return mlirAffineMapGetNumInputs(self); }) 815 .def_property_readonly( 816 "n_symbols", 817 [](PyAffineMap &self) { return mlirAffineMapGetNumSymbols(self); }) 818 .def_property_readonly("results", [](PyAffineMap &self) { 819 return PyAffineMapExprList(self); 820 }); 821 PyAffineMapExprList::bind(m); 822 823 //---------------------------------------------------------------------------- 824 // Mapping of PyIntegerSet. 825 //---------------------------------------------------------------------------- 826 py::class_<PyIntegerSet>(m, "IntegerSet", py::module_local()) 827 .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, 828 &PyIntegerSet::getCapsule) 829 .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyIntegerSet::createFromCapsule) 830 .def("__eq__", [](PyIntegerSet &self, 831 PyIntegerSet &other) { return self == other; }) 832 .def("__eq__", [](PyIntegerSet &self, py::object other) { return false; }) 833 .def("__str__", 834 [](PyIntegerSet &self) { 835 PyPrintAccumulator printAccum; 836 mlirIntegerSetPrint(self, printAccum.getCallback(), 837 printAccum.getUserData()); 838 return printAccum.join(); 839 }) 840 .def("__repr__", 841 [](PyIntegerSet &self) { 842 PyPrintAccumulator printAccum; 843 printAccum.parts.append("IntegerSet("); 844 mlirIntegerSetPrint(self, printAccum.getCallback(), 845 printAccum.getUserData()); 846 printAccum.parts.append(")"); 847 return printAccum.join(); 848 }) 849 .def("__hash__", 850 [](PyIntegerSet &self) { 851 return static_cast<size_t>(llvm::hash_value(self.get().ptr)); 852 }) 853 .def_property_readonly( 854 "context", 855 [](PyIntegerSet &self) { return self.getContext().getObject(); }) 856 .def( 857 "dump", [](PyIntegerSet &self) { mlirIntegerSetDump(self); }, 858 kDumpDocstring) 859 .def_static( 860 "get", 861 [](intptr_t numDims, intptr_t numSymbols, py::list exprs, 862 std::vector<bool> eqFlags, DefaultingPyMlirContext context) { 863 if (exprs.size() != eqFlags.size()) 864 throw py::value_error( 865 "Expected the number of constraints to match " 866 "that of equality flags"); 867 if (exprs.empty()) 868 throw py::value_error("Expected non-empty list of constraints"); 869 870 // Copy over to a SmallVector because std::vector has a 871 // specialization for booleans that packs data and does not 872 // expose a `bool *`. 873 SmallVector<bool, 8> flags(eqFlags.begin(), eqFlags.end()); 874 875 SmallVector<MlirAffineExpr> affineExprs; 876 pyListToVector<PyAffineExpr>(exprs, affineExprs, 877 "attempting to create an IntegerSet"); 878 MlirIntegerSet set = mlirIntegerSetGet( 879 context->get(), numDims, numSymbols, exprs.size(), 880 affineExprs.data(), flags.data()); 881 return PyIntegerSet(context->getRef(), set); 882 }, 883 py::arg("num_dims"), py::arg("num_symbols"), py::arg("exprs"), 884 py::arg("eq_flags"), py::arg("context") = py::none()) 885 .def_static( 886 "get_empty", 887 [](intptr_t numDims, intptr_t numSymbols, 888 DefaultingPyMlirContext context) { 889 MlirIntegerSet set = 890 mlirIntegerSetEmptyGet(context->get(), numDims, numSymbols); 891 return PyIntegerSet(context->getRef(), set); 892 }, 893 py::arg("num_dims"), py::arg("num_symbols"), 894 py::arg("context") = py::none()) 895 .def( 896 "get_replaced", 897 [](PyIntegerSet &self, py::list dimExprs, py::list symbolExprs, 898 intptr_t numResultDims, intptr_t numResultSymbols) { 899 if (static_cast<intptr_t>(dimExprs.size()) != 900 mlirIntegerSetGetNumDims(self)) 901 throw py::value_error( 902 "Expected the number of dimension replacement expressions " 903 "to match that of dimensions"); 904 if (static_cast<intptr_t>(symbolExprs.size()) != 905 mlirIntegerSetGetNumSymbols(self)) 906 throw py::value_error( 907 "Expected the number of symbol replacement expressions " 908 "to match that of symbols"); 909 910 SmallVector<MlirAffineExpr> dimAffineExprs, symbolAffineExprs; 911 pyListToVector<PyAffineExpr>( 912 dimExprs, dimAffineExprs, 913 "attempting to create an IntegerSet by replacing dimensions"); 914 pyListToVector<PyAffineExpr>( 915 symbolExprs, symbolAffineExprs, 916 "attempting to create an IntegerSet by replacing symbols"); 917 MlirIntegerSet set = mlirIntegerSetReplaceGet( 918 self, dimAffineExprs.data(), symbolAffineExprs.data(), 919 numResultDims, numResultSymbols); 920 return PyIntegerSet(self.getContext(), set); 921 }, 922 py::arg("dim_exprs"), py::arg("symbol_exprs"), 923 py::arg("num_result_dims"), py::arg("num_result_symbols")) 924 .def_property_readonly("is_canonical_empty", 925 [](PyIntegerSet &self) { 926 return mlirIntegerSetIsCanonicalEmpty(self); 927 }) 928 .def_property_readonly( 929 "n_dims", 930 [](PyIntegerSet &self) { return mlirIntegerSetGetNumDims(self); }) 931 .def_property_readonly( 932 "n_symbols", 933 [](PyIntegerSet &self) { return mlirIntegerSetGetNumSymbols(self); }) 934 .def_property_readonly( 935 "n_inputs", 936 [](PyIntegerSet &self) { return mlirIntegerSetGetNumInputs(self); }) 937 .def_property_readonly("n_equalities", 938 [](PyIntegerSet &self) { 939 return mlirIntegerSetGetNumEqualities(self); 940 }) 941 .def_property_readonly("n_inequalities", 942 [](PyIntegerSet &self) { 943 return mlirIntegerSetGetNumInequalities(self); 944 }) 945 .def_property_readonly("constraints", [](PyIntegerSet &self) { 946 return PyIntegerSetConstraintList(self); 947 }); 948 PyIntegerSetConstraint::bind(m); 949 PyIntegerSetConstraintList::bind(m); 950 } 951