Add python side Numpy dialect wrapper.

pull/1/head
Stella Laurenzo 2020-05-01 10:38:52 -07:00
parent 23a9ffaabe
commit ba0c96b51a
4 changed files with 58 additions and 0 deletions

View File

@ -0,0 +1,51 @@
# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from npcomp.native.mlir.ir import *
__all__ = [
"load_builtin_module",
]
def load_builtin_module(context = None):
"""Loads a module populated with numpy built-ins.
This is not a long-term solution but overcomes some bootstrapping
issues.
>>> m = load_builtin_module()
>>> op = m.region(0).blocks.front.operations.front
>>> print(op.name)
numpy.generic_ufunc
Args:
context: The MLIRContext to use (None to create a new one).
Returns:
A ModuleOp.
"""
if context is None: context = MLIRContext()
return context.parse_asm(_BUILTIN_MODULE_ASM)
_BUILTIN_MODULE_ASM = r"""
numpy.generic_ufunc @numpy.add (
// CHECK-SAME: overload(%arg0: i32, %arg1: i32) -> i32 {
overload(%arg0: i32, %arg1: i32) -> i32 {
// CHECK: addi
%0 = addi %arg0, %arg1 : i32
numpy.ufunc_return %0 : i32
},
// CHECK: overload(%arg0: f32, %arg1: f32) -> f32 {
overload(%arg0: f32, %arg1: f32) -> f32 {
// CHECK: addf
%0 = addf %arg0, %arg1 : f32
numpy.ufunc_return %0 : f32
}
)
"""
if __name__ == "__main__":
import doctest
doctest.testmod()

View File

View File

@ -28,6 +28,9 @@ static OwningModuleRef parseMLIRModuleFromString(StringRef contents,
//===----------------------------------------------------------------------===//
// Internal only template definitions
// Since it is only legal to use explicit instantiations of templates in
// mlir_ir.h, implementations are kept in this module to keep things scoped
// well for the compiler.
//===----------------------------------------------------------------------===//
template <typename ListTy, typename ItemWrapperTy>
@ -45,6 +48,9 @@ void PyIpListWrapper<ListTy, ItemWrapperTy>::bind(py::module m,
};
py::class_<ThisTy>(m, className)
.def_property_readonly(
"front",
[](ThisTy &self) { return ItemWrapperTy(self.list.front()); })
.def("__len__", [](ThisTy &self) { return self.list.size(); })
.def("__iter__",
[](ThisTy &self) {

View File

@ -7,6 +7,7 @@ import sys
TEST_MODULES = (
"npcomp.mlir_ir_test",
"npcomp.dialect.Numpy",
"npcomp.edsc_test",
"npcomp.tracing.context",
"npcomp.tracing.mlir_trace",