mirror of https://github.com/llvm/torch-mlir
52 lines
1.3 KiB
Python
52 lines
1.3 KiB
Python
# 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()
|