| 141 | }; |
| 142 | |
| 143 | CustomOpImpl::CustomOpImpl(const std::string& op_type, uint32_t version) |
| 144 | : m_version(version), m_op_type(op_type) { |
| 145 | if (m_version != CURRENT_VERSION) { |
| 146 | mgb_log_warn( |
| 147 | "the version of loaded custom op %s is %u, but custom op version " |
| 148 | "of the system is %u\n", |
| 149 | op_type.c_str(), m_version, CURRENT_VERSION); |
| 150 | } |
| 151 | |
| 152 | infer_output_device_func = [](const std::vector<Device>& inputs, const Param&, |
| 153 | std::vector<Device>& outputs) -> void { |
| 154 | static UnImpleWarnLog log_once("output_device_infer", "device", "x86"); |
| 155 | for (size_t i = 0; i < outputs.size(); ++i) { |
| 156 | outputs[i] = inputs.size() > 0 ? inputs[0] : Device("x86"); |
| 157 | } |
| 158 | }; |
| 159 | |
| 160 | infer_output_shape_func = [](const std::vector<Shape>& inputs, const Param&, |
| 161 | std::vector<Shape>& outputs) -> void { |
| 162 | static UnImpleWarnLog log_once("output_shape_infer", "shape", "{1}"); |
| 163 | for (size_t i = 0; i < outputs.size(); ++i) { |
| 164 | outputs[i] = inputs.size() > 0 ? inputs[0] : Shape({1}); |
| 165 | } |
| 166 | }; |
| 167 | |
| 168 | infer_output_dtype_func = [](const std::vector<DType>& inputs, const Param&, |
| 169 | std::vector<DType>& outputs) -> void { |
| 170 | static UnImpleWarnLog log_once("output_dtype_infer", "dtype", "float32"); |
| 171 | for (size_t i = 0; i < outputs.size(); ++i) { |
| 172 | outputs[i] = inputs.size() > 0 ? inputs[0] : DType("float32"); |
| 173 | } |
| 174 | }; |
| 175 | |
| 176 | infer_output_format_func = [](const std::vector<Format>& inputs, const Param&, |
| 177 | std::vector<Format>& outputs) -> void { |
| 178 | for (size_t i = 0; i < outputs.size(); ++i) { |
| 179 | outputs[i] = inputs.size() > 0 ? inputs[0] : Format("default"); |
| 180 | } |
| 181 | }; |
| 182 | |
| 183 | for (const auto& device : Device::legal_devices()) { |
| 184 | compute_funcs[device] = [](const std::vector<Tensor>&, const Param&, |
| 185 | std::vector<Tensor>& outputs, |
| 186 | const RuntimeArgs&) -> void { |
| 187 | auto device = outputs[0].device(); |
| 188 | mgb_assert( |
| 189 | false, |
| 190 | "There is no forward function for your op on device `%s`. " |
| 191 | "Please implement this function and register it.", |
| 192 | device.str().c_str()); |
| 193 | }; |
| 194 | preprocess_funcs[device] = [](const std::vector<Tensor>&, const Param&, |
| 195 | std::vector<Tensor>&, |
| 196 | const RuntimeArgs&) -> void { return; }; |
| 197 | postprocess_funcs[device] = [](const std::vector<Tensor>&, const Param&, |
| 198 | std::vector<Tensor>&, |
| 199 | const RuntimeArgs&) -> void { return; }; |
| 200 | } |