| #include "task_group_priority.h" |
|
|
| void task_group_priority_cpu( |
| int* group, int* priority, bool* value, bool* output, |
| int batch_size, int task_num, int group_num) |
| { |
| auto temp = torch::make_unique<int[]>(group_num); |
| for(int b=0; b<batch_size; b++) |
| { |
| for(int i=0; i<group_num; i++){ |
| temp[i] = std::numeric_limits<int>::max(); |
| } |
| |
| for(int i=0; i<task_num; i++){ |
| if(value[i]){ |
| continue; |
| } |
| int g = group[i]; |
| int p = priority[i]; |
| if(p < temp[g]){ |
| temp[g] = p; |
| } |
| } |
| |
| for(int i=0; i<task_num; i++){ |
| int g = group[i]; |
| output[i] = priority[i]!=temp[g]; |
| } |
| |
| group += task_num; |
| priority += task_num; |
| value += task_num; |
| output += task_num; |
| } |
| }; |
|
|
| auto task_group_priority( |
| const torch::Tensor& group, |
| const torch::Tensor& priority, |
| const torch::Tensor& value) -> torch::Tensor |
| { |
| auto device = group.device(); |
|
|
| const int batch_size = group.size(0); |
| const int task_num = group.size(1); |
| const int group_num = group.max().item<int>() + 1; |
|
|
| const int _group_num = group.min().item<int>(); |
|
|
| GRL_CHECK(group_num <= task_num && _group_num >= 0, "group value error"); |
|
|
| GRL_CHECK_TENSOR(group, device, false, false, batch_size, task_num); |
| GRL_CHECK_TENSOR(priority, device, false, false, batch_size, task_num); |
| GRL_CHECK_TENSOR(value, device, false, false, batch_size, task_num); |
|
|
| auto output = torch::zeros({batch_size, task_num}, torch::dtype(torch::kBool).device(device)); |
|
|
| switch(device.type()) |
| { |
| case torch::kCPU: |
| task_group_priority_cpu(group.data_ptr<int>(), priority.data_ptr<int>(), value.data_ptr<bool>(), |
| output.data_ptr<bool>(), batch_size, task_num, group_num); |
| break; |
| #ifdef CUDA_FOUND |
| case torch::kCUDA: |
| task_group_priority_cuda(group.data_ptr<int>(), priority.data_ptr<int>(), value.data_ptr<bool>(), |
| output.data_ptr<bool>(), batch_size, task_num, group_num, device.index()); |
| break; |
| #endif |
| default: |
| GRL_ERROR("unsupported device: %s", device.str().c_str()); |
| } |
| |
| return output; |
| }; |
|
|
