[BindingsPython] Add useful member methods to LoopBand and LoopBandList

This commit is contained in:
Hanchen Ye 2021-11-03 01:35:20 -05:00
parent ad8dea7101
commit 5fc73d0bb2
2 changed files with 19 additions and 2 deletions

View File

@ -32,9 +32,19 @@ public:
operator MlirAffineLoopBand() const { return get(); }
MlirAffineLoopBand get() const { return {impl.begin(), impl.end()}; }
size_t size() const { return impl.size(); }
PyAffineLoopBand &dunderIter() { return *this; }
MlirOperation dunderNext() {
if (nextIndex >= impl.size())
throw py::stop_iteration();
return impl[nextIndex++];
}
private:
llvm::SmallVector<MlirOperation, 6> impl;
size_t nextIndex = 0;
};
class PyAffineLoopBandList {
@ -49,6 +59,8 @@ public:
}
}
size_t size() const { return impl.size(); }
PyAffineLoopBandList &dunderIter() { return *this; }
PyAffineLoopBand dunderNext() {
@ -121,10 +133,15 @@ PYBIND11_MODULE(_scalehls, m) {
mlirEmitHlsCpp(mod, accum.getCallback(), accum.getUserData());
});
py::class_<PyAffineLoopBand>(m, "LoopBand", py::module_local());
py::class_<PyAffineLoopBand>(m, "LoopBand", py::module_local())
.def_property_readonly("size", &PyAffineLoopBand::size)
.def("__iter__", &PyAffineLoopBand::dunderIter)
.def("__next__", &PyAffineLoopBand::dunderNext);
py::class_<PyAffineLoopBandList>(m, "LoopBandList", py::module_local())
.def(py::init<MlirOperation>(), py::arg("op"),
"Initialize with all loop bands contained by the operation")
.def_property_readonly("size", &PyAffineLoopBandList::size)
.def("__iter__", &PyAffineLoopBandList::dunderIter)
.def("__next__", &PyAffineLoopBandList::dunderNext);
}

View File

@ -50,7 +50,7 @@ def main():
scalehls.apply_affine_loop_perfection(band)
# scalehls.apply_affine_loop_order_opt(band)
scalehls.apply_remove_variable_bound(band)
scalehls.apply_loop_pipelining(band, 2, 3) # pipelineLoc, targetII
scalehls.apply_loop_pipelining(band, band.size - 1, 3) # targetII
scalehls.apply_legalize_to_hlscpp(op.operation, True) # topFunc
scalehls.apply_memory_access_opt(op.operation)
scalehls.apply_array_partition(op.operation)