| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, |
| * software distributed under the License is distributed on an |
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| * KIND, either express or implied. See the License for the |
| * specific language governing permissions and limitations |
| * under the License. |
| */ |
| |
| /*! |
| * \file np_unique_op.cu |
| */ |
| |
| #include "./np_unique_op.h" |
| |
| namespace mxnet { |
| namespace op { |
| |
| struct UniqueComputeAuxGPUKernel { |
| // assume that idx have been flattened to a 1-D tensor (N,) |
| // assume that out_data and in_data have been flattened to 2-D tensors, (N, M) and (K, M) |
| // M is the number of columns of in_data and out_data |
| // i is the index of out_data |
| template <typename DType> |
| MSHADOW_XINLINE static void Map(dim_t i, |
| DType* out_data, |
| const DType* in_data, |
| const dim_t* idx, |
| const dim_t M) { |
| dim_t j = idx[i / M]; |
| out_data[i] = in_data[j * M + i % M]; |
| } |
| }; |
| |
| struct UniqueComputeMaskGPUKernel { |
| template <typename DType> |
| MSHADOW_XINLINE static void Map(dim_t i, |
| dim_t* out_data, |
| const DType* in_data, |
| const dim_t numel) { |
| if (i == 0) { |
| out_data[i] = 1; |
| } else { |
| out_data[i] = 0; |
| for (dim_t j = 0; j < numel; ++j) { |
| if (in_data[(i - 1) * numel + j] != in_data[i * numel + j]) { |
| out_data[i] = 1; |
| break; |
| } |
| } |
| } |
| } |
| }; |
| |
| template <typename DType> |
| struct argsort_functor1d : public thrust::unary_function<dim_t, dim_t> { |
| explicit argsort_functor1d(DType* data_) : data(data_) {} |
| __device__ bool operator()(dim_t a, dim_t b) { |
| return data[a] < data[b]; |
| } |
| DType* data; |
| }; |
| |
| template <typename DType> |
| struct argsort_functor2d : public thrust::unary_function<dim_t, dim_t> { |
| argsort_functor2d(DType* data_, dim_t numel_) : data(data_), numel(numel_) {} |
| __device__ bool operator()(dim_t a, dim_t b) { |
| for (dim_t i = 0; i < numel; ++i) { |
| DType lhs = data[i + a * numel]; |
| DType rhs = data[i + b * numel]; |
| if (lhs < rhs) { |
| return true; |
| } else if (lhs > rhs) { |
| return false; |
| } |
| } |
| return false; |
| } |
| DType* data; |
| dim_t numel; |
| }; |
| |
| void NumpyUniqueGPUNoneAxisImpl(const NumpyUniqueParam& param, |
| const OpContext& ctx, |
| const std::vector<NDArray>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<NDArray>& outputs) { |
| MXNET_NO_FLOAT16_TYPE_SWITCH(outputs[0].dtype(), DType, { |
| mshadow::Stream<gpu>* stream = ctx.get_stream<gpu>(); |
| cudaStream_t cuda_stream = mshadow::Stream<gpu>::GetStream(stream); |
| auto policy = thrust::cuda::par.on(stream->stream_); |
| |
| DType* input_data = inputs[0].data().dptr<DType>(); |
| dim_t input_size = inputs[0].shape().Size(); |
| // argsort, result in perm |
| thrust::device_vector<dim_t> perm(input_size); |
| thrust::sequence(policy, perm.begin(), perm.end()); |
| thrust::stable_sort(policy, perm.begin(), perm.end(), argsort_functor1d<DType>(input_data)); |
| // sorted data in aux |
| thrust::device_vector<DType> aux(input_size); |
| mxnet_op::Kernel<UniqueComputeAuxGPUKernel, gpu>::Launch(stream, |
| input_size, |
| thrust::raw_pointer_cast(aux.data()), |
| input_data, |
| thrust::raw_pointer_cast(perm.data()), |
| 1); |
| // calculate unique mask |
| thrust::device_vector<dim_t> mask(input_size); |
| mxnet_op::Kernel<UniqueComputeMaskGPUKernel, gpu>::Launch(stream, |
| input_size, |
| thrust::raw_pointer_cast(mask.data()), |
| thrust::raw_pointer_cast(aux.data()), |
| 1); |
| // Calculate prefix sum |
| thrust::device_vector<int32_t> prefix_sum(input_size, 0); |
| thrust::inclusive_scan(policy, mask.begin(), mask.end(), prefix_sum.begin()); |
| int32_t valid_num = 0; |
| CUDA_CALL(cudaMemcpyAsync(&valid_num, |
| thrust::raw_pointer_cast(&prefix_sum[input_size - 1]), |
| sizeof(int32_t), |
| cudaMemcpyDeviceToHost, |
| cuda_stream)); |
| CUDA_CALL(cudaStreamSynchronize(cuda_stream)); |
| // set the output shape forcefully |
| mxnet::TShape s(1, valid_num); |
| const_cast<NDArray&>(outputs[0]).Init(s); |
| // launch kernal to obtain unique array, reuse boolean_mask kernel |
| mxnet_op::Kernel<BooleanMaskForwardKernel, gpu>::Launch( |
| stream, |
| input_size, |
| outputs[0].data().dptr<DType>(), |
| thrust::raw_pointer_cast(aux.data()), |
| thrust::raw_pointer_cast(prefix_sum.data()), |
| 1); |
| // handle other optional outputs |
| int output_flag = 0; |
| if (param.return_index) { |
| output_flag += 1; |
| const_cast<NDArray&>(outputs[output_flag]).Init(s); |
| dim_t* unique_indices = outputs[output_flag].data().dptr<dim_t>(); |
| // reuse boolean_mask kernel |
| mxnet_op::Kernel<BooleanMaskForwardKernel, gpu>::Launch( |
| stream, |
| input_size, |
| unique_indices, |
| thrust::raw_pointer_cast(perm.data()), |
| thrust::raw_pointer_cast(prefix_sum.data()), |
| 1); |
| } |
| if (param.return_inverse) { |
| output_flag += 1; |
| const_cast<NDArray&>(outputs[output_flag]).Init(mxnet::TShape(1, input_size)); |
| dim_t* unique_inverse = outputs[output_flag].data().dptr<dim_t>(); |
| mxnet_op::Kernel<UniqueReturnInverseKernel, gpu>::Launch( |
| stream, |
| input_size, |
| unique_inverse, |
| thrust::raw_pointer_cast(prefix_sum.data()), |
| thrust::raw_pointer_cast(perm.data())); |
| } |
| if (param.return_counts) { |
| output_flag += 1; |
| thrust::device_vector<dim_t> idx(valid_num + 1); |
| auto iter = idx.begin(); |
| for (dim_t i = 0; i < input_size; ++i) { |
| if (mask[i]) { |
| *iter = i; |
| ++iter; |
| } |
| } |
| *iter = input_size; |
| const_cast<NDArray&>(outputs[output_flag]).Init(s); |
| dim_t* unique_counts = outputs[output_flag].data().dptr<dim_t>(); |
| mxnet_op::Kernel<UniqueReturnCountsKernel, gpu>::Launch( |
| stream, valid_num, unique_counts, thrust::raw_pointer_cast(idx.data())); |
| } |
| }); |
| } |
| |
| void NumpyUniqueGPUImpl(const NumpyUniqueParam& param, |
| const OpContext& ctx, |
| const std::vector<NDArray>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<NDArray>& outputs) { |
| CHECK(param.axis.value() >= -1 * inputs[0].shape().ndim() && |
| param.axis.value() < inputs[0].shape().ndim()) |
| << "Axis should be in the range of [-r, r-1] where r is the rank of input tensor"; |
| MXNET_NO_FLOAT16_TYPE_SWITCH(outputs[0].dtype(), DType, { |
| using namespace mshadow; |
| using namespace mshadow::expr; |
| Stream<gpu>* stream = ctx.get_stream<gpu>(); |
| cudaStream_t cuda_stream = Stream<gpu>::GetStream(stream); |
| auto policy = thrust::cuda::par.on(stream->stream_); |
| const index_t actual_axis = |
| param.axis.value() + ((param.axis.value() < 0) ? inputs[0].shape().ndim() : 0); |
| // reshape tensor to [origin_shape[axis], -1] |
| const mxnet::TShape origin_shape = inputs[0].shape(); |
| Tensor<gpu, 3, DType> input_tensor_3d = |
| inputs[0].data().FlatTo3D<gpu, DType>(actual_axis, stream); |
| Tensor<gpu, 1, DType> workspace = ctx.requested[0].get_space_typed<gpu, 1, DType>( |
| Shape1(input_tensor_3d.shape_.Size() * 2), stream); |
| Tensor<gpu, 3, DType> input_tensor( |
| workspace.dptr_, |
| Shape3(input_tensor_3d.shape_[1], input_tensor_3d.shape_[0], input_tensor_3d.shape_[2]), |
| stream); |
| input_tensor = swapaxis<1, 0>(input_tensor_3d); |
| const Shape<3> temp_shape = input_tensor.shape_; |
| DType* input_data = input_tensor.dptr_; |
| dim_t numel = temp_shape[1] * temp_shape[2]; |
| // argsort, result in perm |
| thrust::device_vector<dim_t> perm(temp_shape[0]); |
| thrust::sequence(policy, perm.begin(), perm.end()); |
| thrust::stable_sort( |
| policy, perm.begin(), perm.end(), argsort_functor2d<DType>(input_data, numel)); |
| // sorted data in aux |
| Tensor<gpu, 2, DType> aux(workspace.dptr_ + input_tensor_3d.shape_.Size(), |
| Shape2(temp_shape[0], temp_shape[1] * temp_shape[2]), |
| stream); |
| mxnet_op::Kernel<UniqueComputeAuxGPUKernel, gpu>::Launch(stream, |
| temp_shape.Size(), |
| aux.dptr_, |
| input_data, |
| thrust::raw_pointer_cast(perm.data()), |
| numel); |
| // calculate unique mask |
| thrust::device_vector<dim_t> mask(temp_shape[0]); |
| mxnet_op::Kernel<UniqueComputeMaskGPUKernel, gpu>::Launch( |
| stream, temp_shape[0], thrust::raw_pointer_cast(mask.data()), aux.dptr_, numel); |
| // calculate prefix sum |
| thrust::device_vector<int32_t> prefix_sum(temp_shape[0], 0); |
| thrust::inclusive_scan(policy, mask.begin(), mask.end(), prefix_sum.begin()); |
| int32_t valid_num = 0; |
| CUDA_CALL(cudaMemcpyAsync(&valid_num, |
| thrust::raw_pointer_cast(&prefix_sum[temp_shape[0] - 1]), |
| sizeof(int32_t), |
| cudaMemcpyDeviceToHost, |
| cuda_stream)); |
| CUDA_CALL(cudaStreamSynchronize(cuda_stream)); |
| // store the temp output data, reuse the space of 'input_tensor' |
| Tensor<gpu, 3, DType> temp_tensor( |
| workspace.dptr_, Shape3(valid_num, temp_shape[1], temp_shape[2]), stream); |
| // launch kernal to obtain unique array, reuse boolean_mask kernel |
| mxnet_op::Kernel<BooleanMaskForwardKernel, gpu>::Launch( |
| stream, |
| temp_shape.Size(), |
| temp_tensor.dptr_, |
| aux.dptr_, |
| thrust::raw_pointer_cast(prefix_sum.data()), |
| numel); |
| // set the output shape forcefully and swap axis back |
| mxnet::TShape out_shape(origin_shape); |
| out_shape[actual_axis] = valid_num; |
| const_cast<NDArray&>(outputs[0]).Init(out_shape); |
| Tensor<gpu, 3, DType> output_tensor( |
| outputs[0].data().dptr<DType>(), Shape3(temp_shape[1], valid_num, temp_shape[2]), stream); |
| output_tensor = swapaxis<1, 0>(temp_tensor); |
| // handle other optional outputs |
| int output_flag = 0; |
| if (param.return_index) { |
| output_flag += 1; |
| const_cast<NDArray&>(outputs[output_flag]).Init(mxnet::TShape(1, valid_num)); |
| dim_t* unique_indices = outputs[output_flag].data().dptr<dim_t>(); |
| // reuse boolean_mask kernel |
| mxnet_op::Kernel<BooleanMaskForwardKernel, gpu>::Launch( |
| stream, |
| temp_shape[0], |
| unique_indices, |
| thrust::raw_pointer_cast(perm.data()), |
| thrust::raw_pointer_cast(prefix_sum.data()), |
| 1); |
| } |
| if (param.return_inverse) { |
| output_flag += 1; |
| const_cast<NDArray&>(outputs[output_flag]).Init(mxnet::TShape(1, temp_shape[0])); |
| dim_t* unique_inverse = outputs[output_flag].data().dptr<dim_t>(); |
| mxnet_op::Kernel<UniqueReturnInverseKernel, gpu>::Launch( |
| stream, |
| temp_shape[0], |
| unique_inverse, |
| thrust::raw_pointer_cast(prefix_sum.data()), |
| thrust::raw_pointer_cast(perm.data())); |
| } |
| if (param.return_counts) { |
| output_flag += 1; |
| thrust::device_vector<dim_t> idx(valid_num + 1); |
| auto iter = idx.begin(); |
| for (dim_t i = 0; i < temp_shape[0]; ++i) { |
| if (mask[i]) { |
| *iter = i; |
| ++iter; |
| } |
| } |
| *iter = temp_shape[0]; |
| const_cast<NDArray&>(outputs[output_flag]).Init(mxnet::TShape(1, valid_num)); |
| dim_t* unique_counts = outputs[output_flag].data().dptr<dim_t>(); |
| mxnet_op::Kernel<UniqueReturnCountsKernel, gpu>::Launch( |
| stream, valid_num, unique_counts, thrust::raw_pointer_cast(idx.data())); |
| } |
| }); |
| } |
| |
| void NumpyUniqueGPUForward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<NDArray>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<NDArray>& outputs) { |
| CHECK_EQ(inputs.size(), 1U); |
| CHECK(req[0] == kWriteTo || req[0] == kWriteInplace); |
| using namespace mshadow; |
| const NumpyUniqueParam& param = nnvm::get<NumpyUniqueParam>(attrs.parsed); |
| if (inputs[0].shape().ndim() == 0) { |
| CHECK(!param.axis.has_value() || param.axis.value() == -1 || param.axis.value() == 0) |
| << "Axis can only be -1 or 0 for scalor tensor"; |
| Stream<gpu>* s = ctx.get_stream<gpu>(); |
| cudaStream_t stream = Stream<gpu>::GetStream(s); |
| mxnet::TShape shape_1(1, 1); |
| const_cast<NDArray&>(outputs[0]).Init(shape_1); |
| MSHADOW_TYPE_SWITCH(outputs[0].dtype(), DType, { |
| CUDA_CALL(cudaMemcpyAsync(outputs[0].data().dptr<DType>(), |
| inputs[0].data().dptr<DType>(), |
| sizeof(DType), |
| cudaMemcpyDeviceToDevice, |
| stream)); |
| }); |
| int output_flag = 0; |
| if (param.return_index) { |
| output_flag += 1; |
| const_cast<NDArray&>(outputs[output_flag]).Init(shape_1); |
| Tensor<gpu, 1, dim_t> outdata = |
| outputs[output_flag].data().get_with_shape<gpu, 1, dim_t>(Shape1(1), s); |
| ASSIGN_DISPATCH(outdata, OpReqType::kWriteTo, 0); |
| } |
| if (param.return_inverse) { |
| output_flag += 1; |
| const_cast<NDArray&>(outputs[output_flag]).Init(shape_1); |
| Tensor<gpu, 1, dim_t> outdata = |
| outputs[output_flag].data().get_with_shape<gpu, 1, dim_t>(Shape1(1), s); |
| ASSIGN_DISPATCH(outdata, OpReqType::kWriteTo, 0); |
| } |
| if (param.return_counts) { |
| output_flag += 1; |
| const_cast<NDArray&>(outputs[output_flag]).Init(shape_1); |
| Tensor<gpu, 1, dim_t> outdata = |
| outputs[output_flag].data().get_with_shape<gpu, 1, dim_t>(Shape1(1), s); |
| ASSIGN_DISPATCH(outdata, OpReqType::kWriteTo, 1); |
| } |
| } else if (inputs[0].shape().Size() == 0) { |
| // If the input tensor is zero size, only a check on axis is needed |
| if (param.axis.has_value()) { |
| int axis = param.axis.value(); |
| if (axis < 0) |
| axis += inputs[0].shape().ndim(); |
| CHECK(axis >= 0 && axis < inputs[0].shape().ndim()) |
| << "Axis must be within the range of input tensor's dimension"; |
| } |
| // set the shapes of outputs |
| mxnet::TShape shape_0(1, 0); |
| const_cast<NDArray&>(outputs[0]).Init(shape_0); |
| int output_flag = 0; |
| if (param.return_index) { |
| output_flag += 1; |
| const_cast<NDArray&>(outputs[output_flag]).Init(shape_0); |
| } |
| if (param.return_inverse) { |
| output_flag += 1; |
| const_cast<NDArray&>(outputs[output_flag]).Init(shape_0); |
| } |
| if (param.return_counts) { |
| output_flag += 1; |
| const_cast<NDArray&>(outputs[output_flag]).Init(shape_0); |
| } |
| } else { |
| if (!param.axis.has_value()) { |
| NumpyUniqueGPUNoneAxisImpl(param, ctx, inputs, req, outputs); |
| } else { |
| NumpyUniqueGPUImpl(param, ctx, inputs, req, outputs); |
| } |
| } |
| } |
| |
| NNVM_REGISTER_OP(_npi_unique).set_attr<FComputeEx>("FComputeEx<gpu>", NumpyUniqueGPUForward); |
| |
| } // namespace op |
| } // namespace mxnet |
