EDSC: support multi-expression loop bounds

MLIR supports 'for' loops with lower(upper) bound defined by taking a
maximum(minimum) of a list of expressions, but does not have first-class affine
constructs for the maximum(minimum).  All these expressions must have affine
provenance, similarly to a single-expression bound.  Add support for
constructing such loops using EDSC.  The expression factory function is called
`edsc::MaxMinFor` to (1) highlight that the maximum(minimum) operation is
applied to the lower(upper) bound expressions and (2) differentiate it from a
`edsc::For` that creates multiple perfectly nested loops (and should arguably
be called `edsc::ForNest`).

PiperOrigin-RevId: 234785996
This commit is contained in:
Alex Zinenko
2019-02-20 06:54:36 -08:00
committed by jpienaar
parent a2a433652d
commit d055a4e100
8 changed files with 211 additions and 43 deletions

View File

@@ -391,6 +391,15 @@ PYBIND11_MODULE(pybind, m) {
SmallVector<edsc_stmt_t, 8> owning;
return PythonStmt(::For(iv, lb, ub, step, makeCStmts(owning, stmts)));
});
m.def("MaxMinFor", [](PythonExpr iv, const py::list &lbs, const py::list &ubs,
PythonExpr step, const py::list &stmts) {
SmallVector<edsc_expr_t, 8> owningLBs;
SmallVector<edsc_expr_t, 8> owningUBs;
SmallVector<edsc_stmt_t, 8> owningStmts;
return PythonStmt(::MaxMinFor(iv, makeCExprs(owningLBs, lbs),
makeCExprs(owningUBs, ubs), step,
makeCStmts(owningStmts, stmts)));
});
m.def("Select", [](PythonExpr cond, PythonExpr e1, PythonExpr e2) {
return PythonExpr(::Select(cond, e1, e2));
});