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:
@@ -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));
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user