namespace py = pybind11;
+struct PythonAttribute;
struct PythonBindable;
struct PythonExpr;
struct PythonStmt;
return declaration;
}
+ // Create a custom op given its name and arguments.
+ PythonExpr op(const std::string &name, PythonType type,
+ const py::list &arguments, const py::list &successors,
+ py::kwargs attributes);
+
+ // Create an integer attribute.
+ PythonAttribute integerAttr(PythonType type, int64_t value);
+
+ // Create a boolean attribute.
+ PythonAttribute boolAttr(bool value);
+
void compile() {
auto created = mlir::ExecutionEngine::create(module.get());
llvm::handleAllErrors(created.takeError(),
edsc_block_t blk;
};
+struct PythonAttribute {
+ PythonAttribute() : attr(nullptr) {}
+ PythonAttribute(const mlir_attr_t &a) : attr(a) {}
+ PythonAttribute(const PythonAttribute &other) = default;
+ operator mlir_attr_t() { return attr; }
+
+ std::string str() {
+ if (!attr)
+ return "##null attr##";
+
+ std::string res;
+ llvm::raw_string_ostream os(res);
+ Attribute::getFromOpaquePointer(reinterpret_cast<const void *>(attr))
+ .print(os);
+ return res;
+ }
+
+ mlir_attr_t attr;
+};
+
struct PythonIndexed : public edsc_indexed_t {
PythonIndexed(PythonExpr e) : edsc_indexed_t{makeIndexed(e)} {}
PythonIndexed(PythonBindable b) : edsc_indexed_t{makeIndexed(b)} {}
edsc_mlir_emitter_t c_emitter;
};
+template <typename ListTy, typename PythonTy, typename Ty>
+ListTy makeCList(SmallVectorImpl<Ty> &owning, const py::list &list) {
+ for (auto &inp : list) {
+ owning.push_back(Ty{inp.cast<PythonTy>()});
+ }
+ return ListTy{owning.data(), owning.size()};
+}
+
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()};
+ return makeCList<edsc_stmt_list_t, PythonStmt>(owning, stmts);
}
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()};
+ return makeCList<edsc_expr_list_t, PythonExpr>(owning, exprs);
}
static mlir_type_list_t makeCTypes(llvm::SmallVectorImpl<mlir_type_t> &owning,
const py::list &types) {
- for (auto &inp : types) {
- owning.push_back(mlir_type_t{inp.cast<PythonType>()});
- }
- return mlir_type_list_t{owning.data(), owning.size()};
+ return makeCList<mlir_type_list_t, PythonType>(owning, types);
+}
+
+static edsc_block_list_t
+makeCBlocks(llvm::SmallVectorImpl<edsc_block_t> &owning,
+ const py::list &blocks) {
+ return makeCList<edsc_block_list_t, PythonBlock>(owning, blocks);
}
PythonExpr::PythonExpr(const PythonBindable &bindable) : expr{bindable.expr} {}
emitter.emitStmts(StmtBlock(block).getBody());
}
+PythonExpr PythonMLIRModule::op(const std::string &name, PythonType type,
+ const py::list &arguments,
+ const py::list &successors,
+ py::kwargs attributes) {
+ SmallVector<edsc_expr_t, 8> owningExprs;
+ SmallVector<edsc_block_t, 4> owningBlocks;
+ SmallVector<mlir_named_attr_t, 4> owningAttrs;
+ SmallVector<std::string, 4> owningAttrNames;
+
+ owningAttrs.reserve(attributes.size());
+ owningAttrNames.reserve(attributes.size());
+ for (const auto &kvp : attributes) {
+ owningAttrNames.push_back(kvp.first.str());
+ auto value = kvp.second.cast<PythonAttribute>();
+ owningAttrs.push_back({owningAttrNames.back().c_str(), value});
+ }
+
+ return PythonExpr(::Op(mlir_context_t(&mlirContext), name.c_str(), type,
+ makeCExprs(owningExprs, arguments),
+ makeCBlocks(owningBlocks, successors),
+ {owningAttrs.data(), owningAttrs.size()}));
+}
+
+PythonAttribute PythonMLIRModule::integerAttr(PythonType type, int64_t value) {
+ return PythonAttribute(::makeIntegerAttr(type, value));
+}
+
+PythonAttribute PythonMLIRModule::boolAttr(bool value) {
+ return PythonAttribute(::makeBoolAttr(&mlirContext, value));
+}
+
PythonBlock PythonBlock::set(const py::list &stmts) {
SmallVector<edsc_stmt_t, 8> owning;
::BlockSetBody(blk, makeCStmts(owning, stmts));
.def("set", &PythonBlock::set)
.def("__str__", &PythonBlock::str);
+ py::class_<PythonAttribute>(m, "Attribute",
+ "Wrapping class for mlir::Attribute")
+ .def(py::init<PythonAttribute>())
+ .def("__str__", &PythonAttribute::str);
+
py::class_<PythonType>(m, "Type", "Wrapping class for mlir::Type.")
.def(py::init<PythonType>())
.def("__str__", &PythonType::str);
"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("op", &PythonMLIRModule::op, py::arg("name"), py::arg("type"),
+ py::arg("arguments"), py::arg("successors") = py::list(),
+ "Creates a new expression identified by its canonical name.")
+ .def("boolAttr", &PythonMLIRModule::boolAttr,
+ "Creates an mlir::BoolAttr with the given value")
+ .def(
+ "integerAttr", &PythonMLIRModule::integerAttr,
+ "Creates an mlir::IntegerAttr of the given type with the given value "
+ "in the context associated with this MLIR module.")
.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 "
str = expr.__str__()
self.assertIn("($1 * ($2 + $3))", str)
+ def testCustomOp(self):
+ with E.ContextManager():
+ a, b = (E.Expr(E.Bindable(self.i32Type)) for _ in range(2))
+ c1 = self.module.op(
+ "constant",
+ self.i32Type, [],
+ value=self.module.integerAttr(self.i32Type, 42))
+ expr = self.module.op("addi", self.i32Type, [c1, b])
+ str = expr.__str__()
+ self.assertIn("addi(42, $2)", str)
+
def testOneLoop(self):
with E.ContextManager():
i, lb, ub, step = list(
self.module.compile()
self.assertNotEqual(self.module.get_engine_address(), 0)
+ def testCustomOpEmission(self):
+ f = self.module.make_function("fooer", [self.i32Type, self.i32Type], [])
+ with E.ContextManager():
+ emitter = E.MLIRFunctionEmitter(f)
+ funcArg1, funcArg2 = emitter.bind_function_arguments()
+ boolAttr = self.module.boolAttr(True)
+ expr = self.module.op(
+ "foo", self.i32Type, [funcArg1, funcArg2], attr=boolAttr)
+ block = E.Block([E.Stmt(expr), E.Return()])
+ emitter.emit_inplace(block)
+
+ code = str(f)
+ self.assertIn('%0 = "foo"(%arg0, %arg1) {attr: true} : (i32, i32) -> i32',
+ code)
+
+ # Create 'addi' using the generic Op interface. We need an operation known
+ # to the execution engine so that the engine can compile it.
+ def testCustomOpCompilation(self):
+ f = self.module.make_function("adder", [self.i32Type], [])
+ with E.ContextManager():
+ emitter = E.MLIRFunctionEmitter(f)
+ funcArg, = emitter.bind_function_arguments()
+ c1 = self.module.op(
+ "constant",
+ self.i32Type, [],
+ value=self.module.integerAttr(self.i32Type, 42))
+ expr = self.module.op("addi", self.i32Type, [c1, funcArg])
+ block = E.Block([E.Stmt(expr), E.Return()])
+ emitter.emit_inplace(block)
+ self.module.compile()
+ self.assertNotEqual(self.module.get_engine_address(), 0)
+
+
def testMLIREmission(self):
shape = [3, 4, 5]
m = self.module.make_memref_type(self.f32Type, shape)
typedef const void *mlir_type_t;
/// Opaque C type for mlir::Function*.
typedef void *mlir_func_t;
+/// Opaque C type for mlir::Attribute.
+typedef const void *mlir_attr_t;
/// Opaque C type for mlir::edsc::MLIREmiter.
typedef void *edsc_mlir_emitter_t;
/// Opaque C type for mlir::edsc::Expr.
uint64_t n;
} edsc_indexed_list_t;
+typedef struct {
+ edsc_block_t *list;
+ uint64_t n;
+} edsc_block_list_t;
+
+typedef struct {
+ const char *name;
+ mlir_attr_t value;
+} mlir_named_attr_t;
+
+typedef struct {
+ mlir_named_attr_t *list;
+ uint64_t n;
+} mlir_named_attr_list_t;
+
/// Minimal C API for exposing EDSCs to Swift, Python and other languages.
/// Returns a simple scalar mlir::Type using the following convention:
/// Returns an `mlir::IndexType`.
mlir_type_t makeIndexType(mlir_context_t context);
+/// Returns an `mlir::IntegerAttr` of the specified type that contains the given
+/// value.
+mlir_attr_t makeIntegerAttr(mlir_type_t type, int64_t value);
+
+/// Returns an `mlir::BoolAttr` with the given value.
+mlir_attr_t makeBoolAttr(mlir_context_t context, bool value);
+
/// Returns the arity of `function`.
unsigned getFunctionArity(mlir_func_t function);
/// - `indexed` must not have been indexed previously.
edsc_indexed_t index(edsc_indexed_t indexed, edsc_expr_list_t indices);
+/// Returns an opaque expression that will emit an abstract operation identified
+/// by its name.
+edsc_expr_t Op(mlir_context_t context, const char *name, mlir_type_t resultType,
+ edsc_expr_list_t arguments, edsc_block_list_t successors,
+ mlir_named_attr_list_t attrs);
+
/// Returns an opaque expression that will emit an mlir::LoadOp.
edsc_expr_t Load(edsc_indexed_t indexed, edsc_expr_list_t indices);
void print(raw_ostream &os) const;
void dump() const;
+ /// Get an opaque pointer to the attribute.
+ const void *getAsOpaquePointer() const { return attr; }
+ /// Construct an attribute from the opaque pointer representation.
+ static Attribute getFromOpaquePointer(const void *ptr) {
+ return Attribute(
+ const_cast<ImplType *>(reinterpret_cast<const ImplType *>(ptr)));
+ }
+
friend ::llvm::hash_code hash_value(Attribute arg);
protected:
return res;
}
+template <typename Target, size_t N, typename Source>
+SmallVector<Target, N> convertCList(Source list) {
+ SmallVector<Target, N> result;
+ result.reserve(list.n);
+ for (unsigned i = 0; i < list.n; ++i) {
+ result.push_back(Target(list.list[i]));
+ }
+ return result;
+}
+
+SmallVector<StmtBlock, 4> makeBlocks(edsc_block_list_t list) {
+ return convertCList<StmtBlock, 4>(list);
+}
+
static llvm::SmallVector<Expr, 8> makeExprs(edsc_expr_list_t exprList) {
llvm::SmallVector<Expr, 8> exprs;
exprs.reserve(exprList.n);
}
}
+edsc_expr_t Op(mlir_context_t context, const char *name, mlir_type_t resultType,
+ edsc_expr_list_t arguments, edsc_block_list_t successors,
+ mlir_named_attr_list_t attrs) {
+ mlir::MLIRContext *ctx = reinterpret_cast<mlir::MLIRContext *>(context);
+
+ auto blocks = makeBlocks(successors);
+
+ SmallVector<NamedAttribute, 4> attributes;
+ attributes.reserve(attrs.n);
+ for (int i = 0; i < attrs.n; ++i) {
+ auto attribute = Attribute::getFromOpaquePointer(
+ reinterpret_cast<const void *>(attrs.list[i].value));
+ auto name = Identifier::get(attrs.list[i].name, ctx);
+ attributes.emplace_back(name, attribute);
+ }
+
+ return VariadicExpr(
+ name, makeExprs(arguments),
+ Type::getFromOpaquePointer(reinterpret_cast<const void *>(resultType)),
+ attributes, blocks);
+}
+
Expr mlir::edsc::alloc(llvm::ArrayRef<Expr> sizes, Type memrefType) {
return VariadicExpr::make<AllocOp>(sizes, memrefType);
}
// Special case for integer constants that are printed as is. Use
// sign-extended result for everything but i1 (booleans).
if (this->is_op<ConstantIndexOp>() || this->is_op<ConstantIntOp>()) {
+ assert(getAttribute("value"));
APInt value = getAttribute("value").cast<IntegerAttr>().getValue();
if (value.getBitWidth() == 1)
os << value.getZExtValue();
return mlir_type_t{type.getAsOpaquePointer()};
}
+mlir_attr_t makeIntegerAttr(mlir_type_t type, int64_t value) {
+ auto ty = Type::getFromOpaquePointer(reinterpret_cast<const void *>(type));
+ auto attr = IntegerAttr::get(ty, value);
+ return mlir_attr_t{attr.getAsOpaquePointer()};
+}
+
+mlir_attr_t makeBoolAttr(mlir_context_t context, bool value) {
+ auto *ctx = reinterpret_cast<mlir::MLIRContext *>(context);
+ auto attr = BoolAttr::get(value, ctx);
+ return mlir_attr_t{attr.getAsOpaquePointer()};
+}
+
unsigned getFunctionArity(mlir_func_t function) {
auto *f = reinterpret_cast<mlir::Function *>(function);
return f->getNumArguments();
projects / platform / upstream / llvm.git / commitdiff