shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity
llvm-project/mlir/bindings/python/pybind.cpp
Alex Zinenko 59a209721e EDSC: support call instructions 
Introduce support for binding MLIR functions as constant expressions.  Standard
constant operation supports functions as possible constant values.

Provide C APIs to look up existing named functions in an MLIR module and expose
them to the Python bindings.  Provide Python bindings to declare a function in
an MLIR module without defining it and to add a definition given a function
declaration.  These declarations are useful when attempting to link MLIR
modules with, e.g., the standard library.

Introduce EDSC support for direct and indirect calls to other MLIR functions.
Internally, an indirect call is always emitted to leverage existing support for
delayed construction of MLIR Values using EDSC Exprs.  If the expression is
bound to a constant function (looked up or declared beforehand), MLIR constant
folding will be able to replace an indirect call by a direct call.  Currently,
only zero- and one-result functions are supported since we don't have support
for multi-valued expressions in EDSC.

Expose function calling interface to Python bindings on expressions by defining
a `__call__` function accepting a variable number of arguments.

PiperOrigin-RevId: 234959444
2019-03-29 16:36:26 -07:00

719 lines
29 KiB
C++
Raw Blame History

#include "third_party/llvm/llvm/include/llvm/ADT/SmallVector.h"
#include "third_party/llvm/llvm/include/llvm/ADT/StringRef.h"
#include "third_party/llvm/llvm/include/llvm/IR/Module.h"
#include "third_party/llvm/llvm/include/llvm/Support/TargetSelect.h"
#include "third_party/llvm/llvm/include/llvm/Support/raw_ostream.h"
#include <cstddef>
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir-c/Core.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/EDSC/MLIREmitter.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/EDSC/Types.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/ExecutionEngine/ExecutionEngine.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/BuiltinOps.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Module.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Pass/Pass.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Target/LLVMIR.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Transforms/Passes.h"
#include "pybind11/pybind11.h"
#include "pybind11/pytypes.h"
#include "pybind11/stl.h"
#include "mlir/IR/Function.h"
#include "mlir/IR/Types.h"
static bool inited = [] {
llvm::InitializeNativeTarget();
llvm::InitializeNativeTargetAsmPrinter();
return true;
}();
namespace mlir {
namespace edsc {
namespace python {
static std::vector<std::unique_ptr<mlir::Pass>> getDefaultPasses(
const std::vector<const mlir::PassInfo *> &mlirPassInfoList = {}) {
std::vector<std::unique_ptr<mlir::Pass>> passList;
passList.reserve(mlirPassInfoList.size() + 4);
// Run each of the passes that were selected.
for (const auto *passInfo : mlirPassInfoList) {
passList.emplace_back(passInfo->createPass());
}
// Append the extra passes for lowering to MLIR.
passList.emplace_back(mlir::createConstantFoldPass());
passList.emplace_back(mlir::createCSEPass());
passList.emplace_back(mlir::createCanonicalizerPass());
passList.emplace_back(mlir::createLowerAffinePass());
return passList;
}
// Run the passes sequentially on the given module.
// Return `nullptr` immediately if any of the passes fails.
static bool runPasses(const std::vector<std::unique_ptr<mlir::Pass>> &passes,
Module *module) {
for (const auto &pass : passes) {
mlir::PassResult result = pass->runOnModule(module);
if (result == mlir::PassResult::Failure || module->verify()) {
llvm::errs() << "Pass failed\n";
return true;
}
}
return false;
}
namespace py = pybind11;
struct PythonBindable;
struct PythonExpr;
struct PythonStmt;
struct PythonBlock;
struct PythonFunction {
PythonFunction() : function{nullptr} {}
PythonFunction(mlir_func_t f) : function{f} {}
PythonFunction(mlir::Function *f) : function{f} {}
operator mlir_func_t() { return function; }
std::string str() {
mlir::Function *f = reinterpret_cast<mlir::Function *>(function);
std::string res;
llvm::raw_string_ostream os(res);
f->print(os);
return res;
}
// If the function does not yet have an entry block, i.e. if it is a function
// declaration, add the entry block, transforming the declaration into a
// definition. Return true if the block was added, false otherwise.
bool define() {
auto *f = reinterpret_cast<mlir::Function *>(function);
if (!f->getBlocks().empty())
return false;
f->addEntryBlock();
return true;
}
mlir_func_t function;
};
struct PythonType {
PythonType() : type{nullptr} {}
PythonType(mlir_type_t t) : type{t} {}
operator mlir_type_t() { return type; }
std::string str() {
mlir::Type f = mlir::Type::getFromOpaquePointer(type);
std::string res;
llvm::raw_string_ostream os(res);
f.print(os);
return res;
}
mlir_type_t type;
};
/// Trivial C++ wrappers make use of the EDSC C API.
struct PythonMLIRModule {
PythonMLIRModule() : mlirContext(), module(new mlir::Module(&mlirContext)) {}
PythonType makeScalarType(const std::string &mlirElemType,
unsigned bitwidth) {
return ::makeScalarType(mlir_context_t{&mlirContext}, mlirElemType.c_str(),
bitwidth);
}
PythonType makeMemRefType(PythonType elemType, std::vector<int64_t> sizes) {
return ::makeMemRefType(mlir_context_t{&mlirContext}, elemType,
int64_list_t{sizes.data(), sizes.size()});
}
PythonType makeIndexType() {
return ::makeIndexType(mlir_context_t{&mlirContext});
}
// Declare a function with the given name, input and output types but do not
// define it.
PythonFunction declareFunction(const std::string &name,
std::vector<PythonType> &inputTypes,
std::vector<PythonType> &outputTypes) {
std::vector<mlir_type_t> ins(inputTypes.begin(), inputTypes.end());
std::vector<mlir_type_t> outs(outputTypes.begin(), outputTypes.end());
auto funcType = ::makeFunctionType(
mlir_context_t{&mlirContext}, mlir_type_list_t{ins.data(), ins.size()},
mlir_type_list_t{outs.data(), outs.size()});
auto *func = new mlir::Function(
UnknownLoc::get(&mlirContext), name,
mlir::Type::getFromOpaquePointer(funcType).cast<FunctionType>());
module->getFunctions().push_back(func);
return func;
}
// Define a function with the given name, input and output types.
PythonFunction makeFunction(const std::string &name,
std::vector<PythonType> &inputTypes,
std::vector<PythonType> &outputTypes) {
auto declaration = declareFunction(name, inputTypes, outputTypes);
declaration.define();
return declaration;
}
void compile() {
auto created = mlir::ExecutionEngine::create(module.get());
llvm::handleAllErrors(created.takeError(),
[](const llvm::ErrorInfoBase &b) {
b.log(llvm::errs());
assert(false);
});
engine = std::move(*created);
}
std::string getIR() {
std::string res;
llvm::raw_string_ostream os(res);
module->print(os);
return res;
}
uint64_t getEngineAddress() {
assert(engine && "module must be compiled into engine first");
return reinterpret_cast<uint64_t>(reinterpret_cast<void *>(engine.get()));
}
PythonFunction getNamedFunction(const std::string &name) {
return module->getNamedFunction(name);
}
private:
mlir::MLIRContext mlirContext;
// One single module in a python-exposed MLIRContext for now.
std::unique_ptr<mlir::Module> module;
std::unique_ptr<mlir::ExecutionEngine> engine;
};
struct ContextManager {
void enter() { context = new ScopedEDSCContext(); }
void exit(py::object, py::object, py::object) {
delete context;
context = nullptr;
}
mlir::edsc::ScopedEDSCContext *context;
};
struct PythonExpr {
PythonExpr() : expr{nullptr} {}
PythonExpr(const PythonBindable &bindable);
PythonExpr(const edsc_expr_t &expr) : expr{expr} {}
operator edsc_expr_t() { return expr; }
std::string str() {
assert(expr && "unexpected empty expr");
return Expr(*this).str();
}
edsc_expr_t expr;
};
struct PythonBindable : public PythonExpr {
explicit PythonBindable(const PythonType &type)
: PythonExpr(edsc_expr_t{makeBindable(type.type)}) {}
PythonBindable(PythonExpr expr) : PythonExpr(expr) {
assert(Expr(expr).isa<Bindable>() && "Expected Bindable");
}
std::string str() {
assert(expr && "unexpected empty expr");
return Expr(expr).str();
}
};
struct PythonStmt {
PythonStmt() : stmt{nullptr} {}
PythonStmt(const edsc_stmt_t &stmt) : stmt{stmt} {}
PythonStmt(const PythonExpr &e) : stmt{makeStmt(e.expr)} {}
operator edsc_stmt_t() { return stmt; }
std::string str() {
assert(stmt && "unexpected empty stmt");
return Stmt(stmt).str();
}
edsc_stmt_t stmt;
};
struct PythonBlock {
PythonBlock() : blk{nullptr} {}
PythonBlock(const edsc_block_t &other) : blk{other} {}
PythonBlock(const PythonBlock &other) = default;
operator edsc_block_t() { return blk; }
std::string str() {
assert(blk && "unexpected empty block");
return StmtBlock(blk).str();
}
edsc_block_t blk;
};
struct PythonIndexed : public edsc_indexed_t {
PythonIndexed(PythonExpr e) : edsc_indexed_t{makeIndexed(e)} {}
PythonIndexed(PythonBindable b) : edsc_indexed_t{makeIndexed(b)} {}
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)),
currentBuilder(currentFunction),
emitter(&currentBuilder, currentFunction->getLoc()) {}
PythonExpr bindConstantBF16(double value);
PythonExpr bindConstantF16(float value);
PythonExpr bindConstantF32(float value);
PythonExpr bindConstantF64(double value);
PythonExpr bindConstantInt(int64_t value, unsigned bitwidth);
PythonExpr bindConstantIndex(int64_t value);
PythonExpr bindConstantFunction(PythonFunction func);
PythonExpr bindFunctionArgument(unsigned pos);
py::list bindFunctionArguments();
py::list bindFunctionArgumentView(unsigned pos);
py::list bindMemRefShape(PythonExpr boundMemRef);
py::list bindIndexedMemRefShape(PythonIndexed boundMemRef) {
return bindMemRefShape(boundMemRef.base);
}
py::list bindMemRefView(PythonExpr boundMemRef);
py::list bindIndexedMemRefView(PythonIndexed boundMemRef) {
return bindMemRefView(boundMemRef.base);
}
void emit(PythonStmt stmt);
void emitBlock(PythonBlock block);
void emitBlockBody(PythonBlock block);
private:
mlir::Function *currentFunction;
mlir::FuncBuilder currentBuilder;
mlir::edsc::MLIREmitter emitter;
edsc_mlir_emitter_t c_emitter;
};
static edsc_stmt_list_t makeCStmts(llvm::SmallVectorImpl<edsc_stmt_t> &owning,
const py::list &stmts) {
for (auto &inp : stmts) {
owning.push_back(edsc_stmt_t{inp.cast<PythonStmt>()});
}
return edsc_stmt_list_t{owning.data(), owning.size()};
}
static edsc_expr_list_t makeCExprs(llvm::SmallVectorImpl<edsc_expr_t> &owning,
const py::list &exprs) {
for (auto &inp : exprs) {
owning.push_back(edsc_expr_t{inp.cast<PythonExpr>()});
}
return edsc_expr_list_t{owning.data(), owning.size()};
}
PythonExpr::PythonExpr(const PythonBindable &bindable) : expr{bindable.expr} {}
PythonExpr MLIRFunctionEmitter::bindConstantBF16(double value) {
return ::bindConstantBF16(edsc_mlir_emitter_t{&emitter}, value);
}
PythonExpr MLIRFunctionEmitter::bindConstantF16(float value) {
return ::bindConstantF16(edsc_mlir_emitter_t{&emitter}, value);
}
PythonExpr MLIRFunctionEmitter::bindConstantF32(float value) {
return ::bindConstantF32(edsc_mlir_emitter_t{&emitter}, value);
}
PythonExpr MLIRFunctionEmitter::bindConstantF64(double value) {
return ::bindConstantF64(edsc_mlir_emitter_t{&emitter}, value);
}
PythonExpr MLIRFunctionEmitter::bindConstantInt(int64_t value,
unsigned bitwidth) {
return ::bindConstantInt(edsc_mlir_emitter_t{&emitter}, value, bitwidth);
}
PythonExpr MLIRFunctionEmitter::bindConstantIndex(int64_t value) {
return ::bindConstantIndex(edsc_mlir_emitter_t{&emitter}, value);
}
PythonExpr MLIRFunctionEmitter::bindConstantFunction(PythonFunction func) {
return ::bindConstantFunction(edsc_mlir_emitter_t{&emitter}, func);
}
PythonExpr MLIRFunctionEmitter::bindFunctionArgument(unsigned pos) {
return ::bindFunctionArgument(edsc_mlir_emitter_t{&emitter},
mlir_func_t{currentFunction}, pos);
}
PythonExpr getPythonType(edsc_expr_t e) { return PythonExpr(e); }
template <typename T> py::list makePyList(llvm::ArrayRef<T> owningResults) {
py::list res;
for (auto e : owningResults) {
res.append(getPythonType(e));
}
return res;
}
py::list MLIRFunctionEmitter::bindFunctionArguments() {
auto arity = getFunctionArity(mlir_func_t{currentFunction});
llvm::SmallVector<edsc_expr_t, 8> owningResults(arity);
edsc_expr_list_t results{owningResults.data(), owningResults.size()};
::bindFunctionArguments(edsc_mlir_emitter_t{&emitter},
mlir_func_t{currentFunction}, &results);
return makePyList(ArrayRef<edsc_expr_t>{owningResults});
}
py::list MLIRFunctionEmitter::bindMemRefShape(PythonExpr boundMemRef) {
auto rank = getBoundMemRefRank(edsc_mlir_emitter_t{&emitter}, boundMemRef);
llvm::SmallVector<edsc_expr_t, 8> owningShapes(rank);
edsc_expr_list_t resultShapes{owningShapes.data(), owningShapes.size()};
::bindMemRefShape(edsc_mlir_emitter_t{&emitter}, boundMemRef, &resultShapes);
return makePyList(ArrayRef<edsc_expr_t>{owningShapes});
}
py::list MLIRFunctionEmitter::bindMemRefView(PythonExpr boundMemRef) {
auto rank = getBoundMemRefRank(edsc_mlir_emitter_t{&emitter}, boundMemRef);
// Own the PythonExpr for the arg as well as all its dims.
llvm::SmallVector<edsc_expr_t, 8> owningLbs(rank);
llvm::SmallVector<edsc_expr_t, 8> owningUbs(rank);
llvm::SmallVector<edsc_expr_t, 8> owningSteps(rank);
edsc_expr_list_t resultLbs{owningLbs.data(), owningLbs.size()};
edsc_expr_list_t resultUbs{owningUbs.data(), owningUbs.size()};
edsc_expr_list_t resultSteps{owningSteps.data(), owningSteps.size()};
::bindMemRefView(edsc_mlir_emitter_t{&emitter}, boundMemRef, &resultLbs,
&resultUbs, &resultSteps);
py::list res;
res.append(makePyList(ArrayRef<edsc_expr_t>{owningLbs}));
res.append(makePyList(ArrayRef<edsc_expr_t>{owningUbs}));
res.append(makePyList(ArrayRef<edsc_expr_t>{owningSteps}));
return res;
}
void MLIRFunctionEmitter::emit(PythonStmt stmt) {
emitter.emitStmt(Stmt(stmt));
}
void MLIRFunctionEmitter::emitBlock(PythonBlock block) {
emitter.emitBlock(StmtBlock(block));
}
void MLIRFunctionEmitter::emitBlockBody(PythonBlock block) {
emitter.emitStmts(StmtBlock(block).getBody());
}
PythonExpr dispatchCall(py::args args, py::kwargs kwargs) {
assert(args.size() != 0);
llvm::SmallVector<edsc_expr_t, 8> exprs;
exprs.reserve(args.size());
for (auto arg : args) {
exprs.push_back(arg.cast<PythonExpr>());
}
edsc_expr_list_t operands{exprs.data() + 1, exprs.size() - 1};
if (kwargs && kwargs.contains("result")) {
for (const auto &kvp : kwargs) {
if (static_cast<std::string>(kvp.first.str()) == "result")
return ::Call1(exprs.front(), kvp.second.cast<PythonType>(), operands);
}
}
return ::Call0(exprs.front(), operands);
}
PYBIND11_MODULE(pybind, m) {
m.doc() =
"Python bindings for MLIR Embedded Domain-Specific Components (EDSCs)";
m.def("version", []() { return "EDSC Python extensions v0.0"; });
m.def("initContext",
[]() { return static_cast<void *>(new ScopedEDSCContext()); });
m.def("deleteContext",
[](void *ctx) { delete reinterpret_cast<ScopedEDSCContext *>(ctx); });
m.def("Block", [](const py::list &stmts) {
SmallVector<edsc_stmt_t, 8> owning;
return PythonBlock(::Block(makeCStmts(owning, stmts)));
});
m.def("For", [](const py::list &ivs, const py::list &lbs, const py::list &ubs,
const py::list &steps, const py::list &stmts) {
SmallVector<edsc_expr_t, 8> owningIVs;
SmallVector<edsc_expr_t, 8> owningLBs;
SmallVector<edsc_expr_t, 8> owningUBs;
SmallVector<edsc_expr_t, 8> owningSteps;
SmallVector<edsc_stmt_t, 8> owningStmts;
return PythonStmt(
::ForNest(makeCExprs(owningIVs, ivs), makeCExprs(owningLBs, lbs),
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("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));
});
m.def("Return", []() {
return PythonStmt(::Return(edsc_expr_list_t{nullptr, 0}));
});
m.def("Return", [](const py::list &returns) {
SmallVector<edsc_expr_t, 8> owningExprs;
return PythonStmt(::Return(makeCExprs(owningExprs, returns)));
});
m.def("ConstantInteger", [](PythonType type, int64_t value) {
return PythonExpr(::ConstantInteger(type, value));
});
#define DEFINE_PYBIND_BINARY_OP(PYTHON_NAME, C_NAME) \
m.def(PYTHON_NAME, [](PythonExpr e1, PythonExpr e2) { \
return PythonExpr(::C_NAME(e1, e2)); \
});
DEFINE_PYBIND_BINARY_OP("Add", Add);
DEFINE_PYBIND_BINARY_OP("Mul", Mul);
DEFINE_PYBIND_BINARY_OP("Sub", Sub);
// DEFINE_PYBIND_BINARY_OP("Div", Div);
DEFINE_PYBIND_BINARY_OP("LT", LT);
DEFINE_PYBIND_BINARY_OP("LE", LE);
DEFINE_PYBIND_BINARY_OP("GT", GT);
DEFINE_PYBIND_BINARY_OP("GE", GE);
DEFINE_PYBIND_BINARY_OP("EQ", EQ);
DEFINE_PYBIND_BINARY_OP("NE", NE);
DEFINE_PYBIND_BINARY_OP("And", And);
DEFINE_PYBIND_BINARY_OP("Or", Or);
#undef DEFINE_PYBIND_BINARY_OP
#define DEFINE_PYBIND_UNARY_OP(PYTHON_NAME, C_NAME) \
m.def(PYTHON_NAME, [](PythonExpr e1) { return PythonExpr(::C_NAME(e1)); });
DEFINE_PYBIND_UNARY_OP("Negate", Negate);
#undef DEFINE_PYBIND_UNARY_OP
py::class_<PythonFunction>(m, "Function",
"Wrapping class for mlir::Function.")
.def(py::init<PythonFunction>())
.def("__str__", &PythonFunction::str)
.def("define", &PythonFunction::define,
"Adds a body to the function if it does not already have one. "
"Returns true if the body was added");
py::class_<PythonBlock>(m, "StmtBlock",
"Wrapping class for mlir::edsc::StmtBlock")
.def(py::init<PythonBlock>())
.def("__str__", &PythonBlock::str);
py::class_<PythonType>(m, "Type", "Wrapping class for mlir::Type.")
.def(py::init<PythonType>())
.def("__str__", &PythonType::str);
py::class_<PythonMLIRModule>(
m, "MLIRModule",
"An MLIRModule is the abstraction that owns the allocations to support "
"compilation of a single mlir::Module into an ExecutionEngine backed by "
"the LLVM ORC JIT. A typical flow consists in creating an MLIRModule, "
"adding functions, compiling the module to obtain an ExecutionEngine on "
"which named functions may be called. For now the only means to retrieve "
"the ExecutionEngine is by calling `get_engine_address`. This mode of "
"execution is limited to passing the pointer to C++ where the function "
"is called. Extending the API to allow calling JIT compiled functions "
"directly require integration with a tensor library (e.g. numpy). This "
"is left as the prerogative of libraries and frameworks for now.")
.def(py::init<>())
.def("declare_function", &PythonMLIRModule::declareFunction,
"Declares a new mlir::Function in the current mlir::Module. The "
"function has no definition and can be linked to an external "
"library.")
.def("make_function", &PythonMLIRModule::makeFunction,
"Defines a new mlir::Function in the current mlir::Module.")
.def(
"make_scalar_type",
[](PythonMLIRModule &instance, const std::string &type,
unsigned bitwidth) {
return instance.makeScalarType(type, bitwidth);
},
py::arg("type"), py::arg("bitwidth") = 0,
"Returns a scalar mlir::Type using the following convention:\n"
" - makeScalarType(c, \"bf16\") return an "
"`mlir::FloatType::getBF16`\n"
" - makeScalarType(c, \"f16\") return an `mlir::FloatType::getF16`\n"
" - makeScalarType(c, \"f32\") return an `mlir::FloatType::getF32`\n"
" - makeScalarType(c, \"f64\") return an `mlir::FloatType::getF64`\n"
" - makeScalarType(c, \"index\") return an `mlir::IndexType::get`\n"
" - makeScalarType(c, \"i\", bitwidth) return an "
"`mlir::IntegerType::get(bitwidth)`\n\n"
" No other combinations are currently supported.")
.def("make_memref_type", &PythonMLIRModule::makeMemRefType,
"Returns an mlir::MemRefType of an elemental scalar. -1 is used to "
"denote symbolic dimensions in the resulting memref shape.")
.def("make_index_type", &PythonMLIRModule::makeIndexType,
"Returns an mlir::IndexType")
.def("compile", &PythonMLIRModule::compile,
"Compiles the mlir::Module to LLVMIR a creates new opaque "
"ExecutionEngine backed by the ORC JIT.")
.def("get_ir", &PythonMLIRModule::getIR,
"Returns a dump of the MLIR representation of the module. This is "
"used for serde to support out-of-process execution as well as "
"debugging purposes.")
.def("get_engine_address", &PythonMLIRModule::getEngineAddress,
"Returns the address of the compiled ExecutionEngine. This is used "
"for in-process execution.")
.def("__str__", &PythonMLIRModule::getIR,
"Get the string representation of the module");
py::class_<ContextManager>(
m, "ContextManager",
"An EDSC context manager is the memory arena containing all the EDSC "
"allocations.\nUsage:\n\n"
"with E.ContextManager() as _:\n i = E.Expr(E.Bindable())\n ...")
.def(py::init<>())
.def("__enter__", &ContextManager::enter)
.def("__exit__", &ContextManager::exit);
py::class_<MLIRFunctionEmitter>(
m, "MLIRFunctionEmitter",
"An MLIRFunctionEmitter is used to fill an empty function body. This is "
"a staged process:\n"
" 1. create or retrieve an mlir::Function `f` with an empty body;\n"
" 2. make an `MLIRFunctionEmitter(f)` to build the current function;\n"
" 3. create leaf Expr that are either Bindable or already Expr that are"
" bound to constants and function arguments by using methods of "
" `MLIRFunctionEmitter`;\n"
" 4. build the function body using Expr, Indexed and Stmt;\n"
" 5. emit the MLIR to implement the function body.")
.def(py::init<PythonFunction>())
.def("bind_constant_bf16", &MLIRFunctionEmitter::bindConstantBF16)
.def("bind_constant_f16", &MLIRFunctionEmitter::bindConstantF16)
.def("bind_constant_f32", &MLIRFunctionEmitter::bindConstantF32)
.def("bind_constant_f64", &MLIRFunctionEmitter::bindConstantF64)
.def("bind_constant_int", &MLIRFunctionEmitter::bindConstantInt)
.def("bind_constant_index", &MLIRFunctionEmitter::bindConstantIndex)
.def("bind_constant_function", &MLIRFunctionEmitter::bindConstantFunction)
.def("bind_function_argument", &MLIRFunctionEmitter::bindFunctionArgument,
"Returns an Expr that has been bound to a positional argument in "
"the current Function.")
.def("bind_function_arguments",
&MLIRFunctionEmitter::bindFunctionArguments,
"Returns a list of Expr where each Expr has been bound to the "
"corresponding positional argument in the current Function.")
.def("bind_memref_shape", &MLIRFunctionEmitter::bindMemRefShape,
"Returns a list of Expr where each Expr has been bound to the "
"corresponding dimension of the memref.")
.def("bind_memref_view", &MLIRFunctionEmitter::bindMemRefView,
"Returns three lists (lower bound, upper bound and step) of Expr "
"where each triplet of Expr has been bound to the minimal offset, "
"extent and stride of the corresponding dimension of the memref.")
.def("bind_indexed_shape", &MLIRFunctionEmitter::bindIndexedMemRefShape,
"Same as bind_memref_shape but returns a list of `Indexed` that "
"support load and store operations")
.def("bind_indexed_view", &MLIRFunctionEmitter::bindIndexedMemRefView,
"Same as bind_memref_view but returns lists of `Indexed` that "
"support load and store operations")
.def("emit", &MLIRFunctionEmitter::emit,
"Emits the MLIR for the EDSC expressions and statements in the "
"current function body.")
.def("emit", &MLIRFunctionEmitter::emitBlock,
"Emits the MLIR for the EDSC statements into a new block")
.def("emit_inplace", &MLIRFunctionEmitter::emitBlockBody,
"Emits the MLIR for the EDSC statements contained in a EDSC block "
"into the current function body without creating a new block");
py::class_<PythonExpr>(m, "Expr", "Wrapping class for mlir::edsc::Expr")
.def(py::init<PythonBindable>())
.def("__add__", [](PythonExpr e1,
PythonExpr e2) { return PythonExpr(::Add(e1, e2)); })
.def("__sub__", [](PythonExpr e1,
PythonExpr e2) { return PythonExpr(::Sub(e1, e2)); })
.def("__mul__", [](PythonExpr e1,
PythonExpr e2) { return PythonExpr(::Mul(e1, e2)); })
// .def("__div__", [](PythonExpr e1, PythonExpr e2) { return
// PythonExpr(::Div(e1, e2)); })
.def("__lt__", [](PythonExpr e1,
PythonExpr e2) { return PythonExpr(::LT(e1, e2)); })
.def("__le__", [](PythonExpr e1,
PythonExpr e2) { return PythonExpr(::LE(e1, e2)); })
.def("__gt__", [](PythonExpr e1,
PythonExpr e2) { return PythonExpr(::GT(e1, e2)); })
.def("__ge__", [](PythonExpr e1,
PythonExpr e2) { return PythonExpr(::GE(e1, e2)); })
.def("__eq__", [](PythonExpr e1,
PythonExpr e2) { return PythonExpr(::EQ(e1, e2)); })
.def("__ne__", [](PythonExpr e1,
PythonExpr e2) { return PythonExpr(::NE(e1, e2)); })
.def("__and__", [](PythonExpr e1,
PythonExpr e2) { return PythonExpr(::And(e1, e2)); })
.def("__or__", [](PythonExpr e1,
PythonExpr e2) { return PythonExpr(::Or(e1, e2)); })
.def("__invert__", [](PythonExpr e) { return PythonExpr(::Negate(e)); })
.def("__call__", &dispatchCall)
.def("__str__", &PythonExpr::str,
R"DOC(Returns the string value for the Expr)DOC");
py::class_<PythonBindable>(
m, "Bindable",
"Wrapping class for mlir::edsc::Bindable.\nA Bindable is a special Expr "
"that can be bound manually to specific MLIR SSA Values.")
.def(py::init<PythonType>())
.def("__str__", &PythonBindable::str);
py::class_<PythonStmt>(m, "Stmt", "Wrapping class for mlir::edsc::Stmt.")
.def(py::init<PythonExpr>())
.def("__str__", &PythonStmt::str,
R"DOC(Returns the string value for the Expr)DOC");
py::class_<PythonIndexed>(
m, "Indexed",
"Wrapping class for mlir::edsc::Indexed.\nAn Indexed is a wrapper class "
"that support load and store operations.")
.def(py::init<PythonExpr>(), R"DOC(Build from existing Expr)DOC")
.def(py::init<PythonBindable>(), R"DOC(Build from existing Bindable)DOC")
.def(
"load",
[](PythonIndexed &instance, const py::list &indices) {
SmallVector<edsc_expr_t, 8> owning;
return PythonExpr(Load(instance, makeCExprs(owning, indices)));
},
R"DOC(Returns an Expr that loads from an Indexed)DOC")
.def(
"store",
[](PythonIndexed &instance, const py::list &indices,
PythonExpr value) {
SmallVector<edsc_expr_t, 8> owning;
return PythonStmt(
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
} // namespace edsc
} // namespace mlir
Reference in New Issue View Git Blame Copy Permalink