torch-mlir/python/npcomp/dialect/Numpy.py

118 lines
3.8 KiB
Python
Raw Normal View History

2020-05-02 01:38:52 +08:00
# 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
import numpy as np
from npcomp.dialect import Basicpy
from _npcomp.mlir import ir
2020-05-02 01:38:52 +08:00
__all__ = [
"load_builtin_module",
"DialectHelper",
2020-05-02 01:38:52 +08:00
]
class DialectHelper(Basicpy.DialectHelper):
r"""Dialect helper.
>>> c = ir.MLIRContext()
>>> h = DialectHelper(c)
>>> m = c.new_module()
>>> tensor_type = h.tensor_type(h.f32_type)
>>> h.builder.insert_block_start(m.first_block)
>>> f = h.func_op("foobar", h.function_type(
... [tensor_type, tensor_type], [tensor_type]),
... create_entry_block=True)
>>> uf = h.numpy_ufunc_call_op("numpy.add", tensor_type,
... *f.first_block.args)
>>> _ = h.return_op(uf.results)
>>> print(m.to_asm())
<BLANKLINE>
<BLANKLINE>
module {
func @foobar(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
%0 = numpy.ufunc_call @numpy.add(%arg0, %arg1) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
return %0 : tensor<*xf32>
}
}
DenseElementsAttrs:
>>> c.dense_elements_attr(np.asarray([1, 2, 3, 4]))
dense<[1, 2, 3, 4]> : tensor<4xsi64>
>>> c.dense_elements_attr(np.asarray([[1, 2], [3, 4]]))
dense<[[1, 2], [3, 4]]> : tensor<2x2xsi64>
>>> c.dense_elements_attr(np.asarray([[1., 2.], [3., 4.]]))
dense<[[1.000000e+00, 2.000000e+00], [3.000000e+00, 4.000000e+00]]> : tensor<2x2xf64>
>>> c.dense_elements_attr(np.asarray([[1., 2.], [3., 4.]], dtype=np.float32))
dense<[[1.000000e+00, 2.000000e+00], [3.000000e+00, 4.000000e+00]]> : tensor<2x2xf32>
Types:
>>> t = DialectHelper(ir.MLIRContext())
>>> t.numpy_any_dtype
!numpy.any_dtype
>>> t.tensor_type(t.numpy_any_dtype, [1, 2, 3])
tensor<1x2x3x!numpy.any_dtype>
>>> t.tensor_type(t.numpy_any_dtype)
tensor<*x!numpy.any_dtype>
>>> t.tensor_type(t.numpy_any_dtype, [-1, 2])
tensor<?x2x!numpy.any_dtype>
>>> t.tensor_type(t.f32_type)
tensor<*xf32>
>>> t.function_type([t.i32_type], [t.f32_type])
(i32) -> f32
>>> t.unknown_array_type
tensor<*x!numpy.any_dtype>
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.numpy_any_dtype = self.context.parse_type("!numpy.any_dtype")
self.unknown_array_type = self.tensor_type(self.numpy_any_dtype)
def numpy_ufunc_call_op(self, callee_symbol, result_type, *args):
"""Creates a numpy.ufunc_call op."""
c = self.context
attrs = c.dictionary_attr(
{"ufunc_ref": c.flat_symbol_ref_attr(callee_symbol)})
return self.op("numpy.ufunc_call", [result_type], args, attrs)
def numpy_narrow_op(self, result_type, operand):
"""Creates a numpy.narrow op."""
return self.op("numpy.narrow", [result_type], [operand])
def numpy_get_slice_op(self, result_type, array, *slice_elements):
return self.op("numpy.get_slice", [result_type],
[array] + list(slice_elements))
def load_builtin_module(context=None):
"""Loads a module populated with numpy built-ins.
2020-05-02 01:38:52 +08:00
This is not a long-term solution but overcomes some bootstrapping
issues.
>>> m = load_builtin_module()
>>> op = m.region(0).blocks.front.operations.front
>>> op.is_registered
True
>>> op.name
'numpy.builtin_ufunc'
2020-05-02 01:38:52 +08:00
Args:
context: The MLIRContext to use (None to create a new one).
Returns:
A ModuleOp.
"""
if context is None: context = ir.MLIRContext()
return context.parse_asm(_BUILTIN_MODULE_ASM)
2020-05-02 01:38:52 +08:00
_BUILTIN_MODULE_ASM = r"""
numpy.builtin_ufunc @numpy.add
numpy.builtin_ufunc @numpy.multiply
2020-05-02 01:38:52 +08:00
"""
if __name__ == "__main__":
import doctest
doctest.testmod()