1 /** 2 * Copyright 2020 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #ifndef MINDSPORE_CCSRC_UTILS_TENSOR_PY_H_ 18 #define MINDSPORE_CCSRC_UTILS_TENSOR_PY_H_ 19 20 #include <memory> 21 #include <string> 22 #include <vector> 23 24 #include "pybind11/numpy.h" 25 26 #include "ir/tensor.h" 27 #include "include/common/np_dtype/np_dtypes.h" 28 29 namespace py = pybind11; 30 31 namespace pybind11 { 32 namespace detail { 33 // Similar to enums in `pybind11/numpy.h`. Determined by doing: 34 // python3 -c 'import numpy as np; print(np.dtype(np.float16).num)' 35 constexpr int NPY_FLOAT16 = 23; 36 37 template <typename T> 38 struct npy_scalar_caster { 39 PYBIND11_TYPE_CASTER(T, _("PleaseOverride")); 40 using Array = array_t<T>; 41 loadnpy_scalar_caster42 bool load(handle src, bool convert) const { 43 // Taken from Eigen casters. Permits either scalar dtype or scalar array. 44 handle type = dtype::of<T>().attr("type"); 45 if (!convert && !isinstance<Array>(src) && !isinstance(src, type)) { 46 return false; 47 } 48 49 Array tmp = Array::ensure(src); 50 if (tmp && tmp.size() == 1 && tmp.ndim() == 0) { 51 this->value = *tmp.data(); 52 return true; 53 } 54 55 return false; 56 } 57 castnpy_scalar_caster58 static handle cast(T src, return_value_policy, handle) { 59 Array tmp({1}); 60 tmp.mutable_at(0) = src; 61 tmp.resize({}); 62 63 // You could also just return the array if you want a scalar array. 64 object scalar = tmp[tuple()]; 65 return scalar.release(); 66 } 67 }; 68 69 template <> 70 struct npy_format_descriptor<float16> { 71 static constexpr auto name = "float16"; 72 static pybind11::dtype dtype() { 73 handle ptr = npy_api::get().PyArray_DescrFromType_(NPY_FLOAT16); 74 return reinterpret_borrow<pybind11::dtype>(ptr); 75 } 76 virtual ~npy_format_descriptor<float16>() {} 77 }; 78 79 template <> 80 struct type_caster<float16> : public npy_scalar_caster<float16> { 81 static constexpr auto name = "float16"; 82 }; 83 84 template <> 85 struct npy_format_descriptor<bfloat16> { 86 static constexpr auto name = "bfloat16"; 87 static pybind11::dtype dtype() { 88 handle ptr = npy_api::get().PyArray_DescrFromType_(mindspore::GetBFloat16NpDType()); 89 return reinterpret_borrow<pybind11::dtype>(ptr); 90 } 91 virtual ~npy_format_descriptor<bfloat16>() {} 92 }; 93 94 template <> 95 struct type_caster<bfloat16> : public npy_scalar_caster<bfloat16> { 96 static constexpr auto name = "bfloat16"; 97 }; 98 } // namespace detail 99 } // namespace pybind11 100 101 // brief mindspore namespace. 102 // 103 // mindspore namespace is the top level namespace of Mindsporeession project. 104 // Other namespace should be a sub namespace of mindspore namespace in the ME project. 105 namespace mindspore { 106 // brief mindspore::tensor namespace 107 // 108 // A sub namespace in ME to support tensor related definition. 109 namespace tensor { 110 // Tensor python wrapper and adapter class. 111 class TensorPy { 112 public: 113 // brief Create Tensor from a numpy array object. 114 // 115 // param input [py::array] Data value of the tensor. 116 // param data_type [TypeId] Data type of the tensor. 117 static TensorPtr MakeTensor(const py::array &input, const TypePtr &type_ptr = nullptr); 118 119 // brief Create Tensor from a numpy array without copy. 120 // 121 // param input [py::array] Data value of the tensor. 122 static TensorPtr MakeTensorOfNumpy(const py::array &input); 123 124 // brief Create Tensor from a numpy array without copy, use persistent tensor data. 125 // 126 // param input [py::array] Data value of the tensor. 127 // param input [py::int_] slice num of data. 128 static TensorPtr MakePersistentDataTensorOfNumpy(const py::array &input, const py::int_ slice_num); 129 130 static py::bytes GetBytes(const Tensor &tensor); 131 132 static py::buffer_info GetPyBufferFromPyArray(const py::array &input); 133 134 static TensorPtr ConvertBytesToTensor(const py::bytes &bytes_obj, const py::tuple &dims, 135 const TypePtr &type_ptr = nullptr); 136 137 static py::array SyncAsNumpy(const Tensor &tensor); 138 139 static py::array AsNumpy(const Tensor &tensor); 140 141 // brief Get slice data as numpy of tensor which use persistent tensor data. 142 // 143 // return [py::array] Slice Data of the tensor at slice_index. 144 static py::array AsNumpyOfSlice(const Tensor &tensor, const int32_t param_key, int slice_index); 145 146 static py::tuple GetPyTupleShape(const Tensor &tensor); 147 148 static py::tuple GetPyTupleStrides(const Tensor &tensor); 149 150 static py::int_ GetPyItemSize(const Tensor &tensor); 151 152 static py::int_ GetPyNBytes(const Tensor &tensor); 153 154 static void FlushFromCache(const Tensor &tensor); 155 156 static void Offload(const Tensor &tensor); 157 // move tensor from device to host, or host to device asynchronously 158 static TensorPtr MoveTo(const Tensor &self, const std::string &to, bool blocking = True); 159 160 static void SetDeviceAddress(const Tensor &tensor, uintptr_t addr, const ShapeVector &shape, const TypePtr type_ptr); 161 }; 162 163 // CSRTensor python wrapper and adapter class. 164 class CSRTensorPy { 165 public: 166 static py::tuple GetPyTupleShape(const CSRTensor &csr_tensor); 167 }; 168 169 // COOTensor python wrapper and adapter class. 170 class COOTensorPy { 171 public: 172 static py::tuple GetPyTupleShape(const COOTensor &coo_tensor); 173 }; 174 175 // RowTensor python wrapper and adapter class. 176 class RowTensorPy { 177 public: 178 static py::tuple GetPyTupleShape(const RowTensor &row_tensor); 179 }; 180 } // namespace tensor 181 } // namespace mindspore 182 183 #endif // MINDSPORE_CCSRC_UTILS_TENSOR_PY_H_ 184