| 144 | } |
| 145 | |
| 146 | void HalideExecutable::invoke( |
| 147 | void* user_context, const FunctionHandle& handle, const VarNodeArray& inputs, |
| 148 | VarNode* output) { |
| 149 | mgb_assert(handle.execute && handle.get_device_interface); |
| 150 | halide_device_interface_t* device_interface = handle.get_device_interface(); |
| 151 | |
| 152 | size_t nr_inputs = m_value_inputs.size(), argv_idx = 0; |
| 153 | void* argv[nr_inputs + 2]; |
| 154 | |
| 155 | halide_buffer_t image_args[nr_inputs + 1]; |
| 156 | size_t nr_dims = (nr_inputs + 1) * TensorLayout::MAX_NDIM; |
| 157 | halide_dimension_t image_dims_buf[nr_dims]; |
| 158 | memset(image_dims_buf, 0, sizeof(halide_dimension_t) * nr_dims); |
| 159 | size_t image_arg_idx = 0; |
| 160 | halide_dimension_t* image_dims_ptr = image_dims_buf; |
| 161 | |
| 162 | auto add_tensor_arg = [&](const DeviceTensorND& tensor) { |
| 163 | int ndim = tensor.layout().ndim; |
| 164 | for (int i = ndim - 1; i >= 0; i--) { |
| 165 | image_dims_ptr->extent = tensor.layout()[i]; |
| 166 | image_dims_ptr->stride = tensor.layout().stride[i]; |
| 167 | image_dims_ptr++; |
| 168 | } |
| 169 | auto dtype = tensor.dtype(); |
| 170 | halide_type_t type = dtype_mgb2halide(dtype); |
| 171 | image_args[image_arg_idx] = { |
| 172 | reinterpret_cast<uint64_t>(tensor.raw_ptr()), |
| 173 | device_interface, |
| 174 | nullptr, |
| 175 | 0, |
| 176 | type, |
| 177 | ndim, |
| 178 | image_dims_ptr - ndim, |
| 179 | nullptr}; |
| 180 | argv[argv_idx++] = &image_args[image_arg_idx++]; |
| 181 | }; |
| 182 | |
| 183 | argv[argv_idx++] = &user_context; |
| 184 | for (auto&& i : m_value_inputs) { |
| 185 | add_tensor_arg(inputs.at(i.first)->dev_tensor()); |
| 186 | } |
| 187 | add_tensor_arg(output->dev_tensor()); |
| 188 | mgb_assert(argv_idx == nr_inputs + 2); |
| 189 | mgb_assert(image_dims_ptr <= image_dims_buf + nr_dims); |
| 190 | auto err = handle.execute(argv); |
| 191 | mgb_throw_if(err, SystemError, "failed to execute halide function: err=%d", err); |
| 192 | } |
| 193 | |
| 194 | halide_type_t HalideExecutable::dtype_mgb2halide(DType dtype) { |
| 195 | if (dtype == dtype::Float32()) { |