mirror of https://github.com/llvm/torch-mlir
Add python side Numpy dialect wrapper.
parent
23a9ffaabe
commit
ba0c96b51a
|
@ -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()
|
|
@ -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) {
|
||||
|
|
|
@ -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",
|
||||
|
|
Loading…
Reference in New Issue