mirror of https://github.com/llvm/torch-mlir
134 lines
4.1 KiB
Python
134 lines
4.1 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
|
|
|
|
import numpy as np
|
|
from npcomp.dialect import Basicpy
|
|
from _npcomp.mlir import ir
|
|
|
|
__all__ = [
|
|
"load_builtin_module",
|
|
"DialectHelper",
|
|
]
|
|
|
|
|
|
class DialectHelper(Basicpy.DialectHelper):
|
|
r"""Dialect helper.
|
|
|
|
>>> c = ir.MLIRContext()
|
|
>>> h = DialectHelper(c, ir.OpBuilder(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:
|
|
>>> c = ir.MLIRContext()
|
|
>>> t = DialectHelper(c, ir.OpBuilder(c))
|
|
>>> t.numpy_any_dtype
|
|
!basicpy.UnknownType
|
|
>>> t.tensor_type(t.numpy_any_dtype, [1, 2, 3])
|
|
tensor<1x2x3x!basicpy.UnknownType>
|
|
>>> t.tensor_type(t.numpy_any_dtype)
|
|
tensor<*x!basicpy.UnknownType>
|
|
>>> t.tensor_type(t.numpy_any_dtype, [-1, 2])
|
|
tensor<?x2x!basicpy.UnknownType>
|
|
>>> t.tensor_type(t.f32_type)
|
|
tensor<*xf32>
|
|
>>> t.function_type([t.i32_type], [t.f32_type])
|
|
(i32) -> f32
|
|
>>> t.unknown_tensor_type
|
|
tensor<*x!basicpy.UnknownType>
|
|
|
|
"""
|
|
|
|
@property
|
|
def numpy_any_dtype(self):
|
|
return self.basicpy_UnknownType
|
|
|
|
@property
|
|
def numpy_unknown_tensor_type(self):
|
|
return self.tensor_type(self.basicpy_UnknownType)
|
|
|
|
@property
|
|
def unknown_array_type(self):
|
|
return self.numpy_NdArrayType(self.basicpy_UnknownType)
|
|
|
|
def numpy_builtin_ufunc_call_op(self, *args, qualified_name, result_type):
|
|
"""Creates a numpy.builtin_ufunc_call op."""
|
|
c = self.context
|
|
attrs = c.dictionary_attr({"qualified_name": c.string_attr(qualified_name)})
|
|
return self.op("numpy.builtin_ufunc_call", [result_type], args, attrs)
|
|
|
|
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.
|
|
|
|
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'
|
|
|
|
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)
|
|
|
|
|
|
_BUILTIN_MODULE_ASM = r"""
|
|
numpy.builtin_ufunc @numpy.add
|
|
numpy.builtin_ufunc @numpy.multiply
|
|
"""
|
|
|
|
if __name__ == "__main__":
|
|
import doctest
|
|
doctest.testmod()
|