caio/clang-p2996
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

EDSC: introduce min/max only usable inside for upper/lower bounds of a loop

Browse Source 
Introduce a type-safe way of building a 'for' loop with max/min bounds in EDSC.
Define new types MaxExpr and MinExpr in C++ EDSC API and expose them to Python
bindings.  Use values of these type to construct 'for' loops with max/min in
newly introduced overloads of the `edsc::For` factory function.  Note that in C
APIs, we still must expose MaxMinFor as a different function because C has no
overloads.  Also note that MaxExpr and MinExpr do _not_ derive from Expr
because they are not allowed to be used in a regular Expr context (which may
produce `affine.apply` instructions not expecting `min` or `max`).

Factory functions `Min` and `Max` in Python can be further overloaded to
produce chains of comparisons and selects on non-index types.  This is not
trivial in C++ since overloaded functions cannot differ by the return type only
(`MaxExpr` or `Expr`) and making `MaxExpr` derive from `Expr` defies the
purpose of type-safe construction.

PiperOrigin-RevId: 234786131
This commit is contained in:
Alex Zinenko
2019-02-20 06:54:53 -08:00
committed by jpienaar
parent d055a4e100
commit 21bd4540f3
5 changed files with 128 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
mlir/bindings/python/pybind.cpp
View File

@@ -222,6 +222,22 @@ struct PythonIndexed : public edsc_indexed_t {
operator PythonExpr() { return PythonExpr(base); }
};
struct PythonMaxExpr {
PythonMaxExpr() : expr(nullptr) {}
PythonMaxExpr(const edsc_max_expr_t &e) : expr(e) {}
operator edsc_max_expr_t() { return expr; }
edsc_max_expr_t expr;
};
struct PythonMinExpr {
PythonMinExpr() : expr(nullptr) {}
PythonMinExpr(const edsc_min_expr_t &e) : expr(e) {}
operator edsc_min_expr_t() { return expr; }
edsc_min_expr_t expr;
};
struct MLIRFunctionEmitter {
MLIRFunctionEmitter(PythonFunction f)
: currentFunction(reinterpret_cast<mlir::Function *>(f.function)),
@@ -386,19 +402,23 @@ PYBIND11_MODULE(pybind, m) {
makeCExprs(owningUBs, ubs), makeCExprs(owningSteps, steps),
makeCStmts(owningStmts, stmts)));
});
m.def("Max", [](const py::list &args) {
SmallVector<edsc_expr_t, 8> owning;
return PythonMaxExpr(::Max(makeCExprs(owning, args)));
});
m.def("Min", [](const py::list &args) {
SmallVector<edsc_expr_t, 8> owning;
return PythonMinExpr(::Min(makeCExprs(owning, args)));
});
m.def("For", [](PythonExpr iv, PythonExpr lb, PythonExpr ub, PythonExpr step,
const py::list &stmts) {
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("For", [](PythonExpr iv, PythonMaxExpr lb, PythonMinExpr ub,
PythonExpr step, const py::list &stmts) {
SmallVector<edsc_stmt_t, 8> owning;
return PythonStmt(::MaxMinFor(iv, lb, ub, step, makeCStmts(owning, stmts)));
});
m.def("Select", [](PythonExpr cond, PythonExpr e1, PythonExpr e2) {
return PythonExpr(::Select(cond, e1, e2));
@@ -623,6 +643,11 @@ PYBIND11_MODULE(pybind, m) {
Store(value, instance, makeCExprs(owning, indices)));
},
R"DOC(Returns the Stmt that stores into an Indexed)DOC");
py::class_<PythonMaxExpr>(m, "MaxExpr",
"Wrapping class for mlir::edsc::MaxExpr");
py::class_<PythonMinExpr>(m, "MinExpr",
"Wrapping class for mlir::edsc::MinExpr");
}
} // namespace python