diff --git a/include/npcomp/Dialect/Numpy/IR/NumpyOps.td b/include/npcomp/Dialect/Numpy/IR/NumpyOps.td index 3b23ffeac..d65d66f49 100644 --- a/include/npcomp/Dialect/Numpy/IR/NumpyOps.td +++ b/include/npcomp/Dialect/Numpy/IR/NumpyOps.td @@ -38,7 +38,8 @@ def Numpy_NarrowOp : Numpy_Op<"narrow", []> { // NdArray type handling //----------------------------------------------------------------------------// -def Numpy_CopyToArray : Numpy_Op<"create_array_from_tensor", [NoSideEffect]> { +def Numpy_CreateArrayFromTensorOp : Numpy_Op<"create_array_from_tensor", [ + NoSideEffect]> { let summary = "Creates an ndarray from a tensor."; let description = [{ Creates a new ndarray that will contain the data of the given tensor. @@ -54,7 +55,7 @@ def Numpy_CopyToArray : Numpy_Op<"create_array_from_tensor", [NoSideEffect]> { }]; } -def Numpy_CopyToTensor : Numpy_Op<"copy_to_tensor", []> { +def Numpy_CopyToTensorOp : Numpy_Op<"copy_to_tensor", []> { let summary = "Copies an ndarray, yielding a value-typed tensor."; let description = [{ The semantics of this operation connote a copy of the data in the source diff --git a/lib/Python/NpcompDialect.cpp b/lib/Python/NpcompDialect.cpp index f70ce8d2e..4e45a8f28 100644 --- a/lib/Python/NpcompDialect.cpp +++ b/lib/Python/NpcompDialect.cpp @@ -11,6 +11,8 @@ #include "npcomp/Dialect/Basicpy/IR/BasicpyDialect.h" #include "npcomp/Dialect/Basicpy/IR/BasicpyOps.h" +#include "npcomp/Dialect/Numpy/IR/NumpyDialect.h" +#include "npcomp/Dialect/Numpy/IR/NumpyOps.h" namespace mlir { namespace NPCOMP { @@ -102,6 +104,26 @@ public: auto op = opBuilder.create( loc, resultType, slotObject, indexAttr); return op.getOperation(); + }) + .def("numpy_create_array_from_tensor_op", + [](BasicpyDialectHelper &self, PyValue source) -> PyOperationRef { + auto sourceType = source.value.getType().dyn_cast(); + if (!sourceType) { + throw py::raiseValueError("expected tensor type for " + "numpy_create_array_from_tensor_op"); + } + auto dtype = sourceType.getElementType(); + auto ndarrayType = + Numpy::NdArrayType::get(dtype, self.getContext()); + OpBuilder &opBuilder = self.pyOpBuilder.getBuilder(true); + Location loc = self.pyOpBuilder.getCurrentLoc(); + auto op = opBuilder.create( + loc, ndarrayType, source.value); + return op.getOperation(); + }) + .def("numpy_NdArrayType", + [](BasicpyDialectHelper &self, PyType dtype) -> PyType { + return Numpy::NdArrayType::get(dtype.type, self.getContext()); }); } }; diff --git a/pytest/NumpyCompiler/array_basics.py b/pytest/NumpyCompiler/array_basics.py new file mode 100644 index 000000000..807e5aea8 --- /dev/null +++ b/pytest/NumpyCompiler/array_basics.py @@ -0,0 +1,18 @@ +# RUN: %PYTHON %s | npcomp-opt -split-input-file | FileCheck %s --dump-input=fail + +import numpy as np +from npcomp.compiler import test_config + +import_global = test_config.create_import_dump_decorator() + +global_data = (np.zeros((2, 3)) + [1.0, 2.0, 3.0] * np.reshape([1.0, 2.0], + (2, 1))) + + +# CHECK-LABEL: func @global_array_to_const +@import_global +def global_array_to_const(): + # CHECK: %[[CST:.*]] = constant dense<{{\[\[}}1.000000e+00, 2.000000e+00, 3.000000e+00], [2.000000e+00, 4.000000e+00, 6.000000e+00]]> : tensor<2x3xf64> + # CHECK: numpy.create_array_from_tensor %[[CST]] : (tensor<2x3xf64>) -> !numpy.ndarray + local_data = global_data + return local_data diff --git a/python/npcomp/compiler/interfaces.py b/python/npcomp/compiler/interfaces.py index 31f83cc83..99de67848 100644 --- a/python/npcomp/compiler/interfaces.py +++ b/python/npcomp/compiler/interfaces.py @@ -113,6 +113,9 @@ class ValueCoderChain(ValueCoder): def __init__(self, sub_coders: Sequence[ValueCoder]): self._sub_coders = tuple(sub_coders) + def __repr__(self): + return "ValueCoderChain({})".format(self._sub_coders) + def code_py_value_as_const(self, env: "Environment", py_value) -> Union[_NotImplementedType, ir.Value]: for sc in self._sub_coders: diff --git a/python/npcomp/compiler/test_config.py b/python/npcomp/compiler/test_config.py index 4c4d575b0..d63c85cb0 100644 --- a/python/npcomp/compiler/test_config.py +++ b/python/npcomp/compiler/test_config.py @@ -11,6 +11,7 @@ from .interfaces import * from .partial_eval_base import * from .target import * from .value_coder_base import * +from .value_coder_numpy import * def create_import_dump_decorator(*, @@ -30,7 +31,10 @@ def create_import_dump_decorator(*, def create_test_config(target_factory: TargetFactory = GenericTarget64): - value_coder = BuiltinsValueCoder() + value_coder = ValueCoderChain([ + BuiltinsValueCoder(), + CreateNumpyValueCoder(), + ]) pe_hook = build_default_partial_eval_hook() return Configuration(target_factory=target_factory, diff --git a/python/npcomp/compiler/value_coder_numpy.py b/python/npcomp/compiler/value_coder_numpy.py new file mode 100644 index 000000000..d3c8e3d2a --- /dev/null +++ b/python/npcomp/compiler/value_coder_numpy.py @@ -0,0 +1,39 @@ +# 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 +"""Value coders for Numpy types.""" + +import numpy as np +from typing import Union + +from _npcomp.mlir import ir + +from . import logging +from .interfaces import * + +__all__ = [ + "CreateNumpyValueCoder", +] + +_NotImplementedType = type(NotImplemented) + + +class NdArrayValueCoder(ValueCoder): + """Value coder for numpy types.""" + __slots__ = [] + + def code_py_value_as_const(self, env: Environment, + py_value) -> Union[_NotImplementedType, ir.Value]: + # TODO: Query for ndarray compat (for duck typed and such) + # TODO: Have a higher level name resolution signal which indicates const + ir_h = env.ir_h + if isinstance(py_value, np.ndarray): + dense_attr = ir_h.context.dense_elements_attr(py_value) + tensor_type = dense_attr.type + tensor_value = ir_h.constant_op(tensor_type, dense_attr).result + return ir_h.numpy_create_array_from_tensor_op(tensor_value).result + return NotImplemented + + +def CreateNumpyValueCoder() -> ValueCoder: + return ValueCoderChain((NdArrayValueCoder(),)) diff --git a/python/npcomp/dialect/Numpy.py b/python/npcomp/dialect/Numpy.py index 69f936088..cf57cfb32 100644 --- a/python/npcomp/dialect/Numpy.py +++ b/python/npcomp/dialect/Numpy.py @@ -14,14 +14,14 @@ __all__ = [ 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]), + ... [tensor_type, tensor_type], [tensor_type]), ... create_entry_block=True) >>> uf = h.numpy_ufunc_call_op("numpy.add", tensor_type, ... *f.first_block.args)