[add] tree_filter kernel support fp16

Author: StevenGrove

furnace/kernels/lib_tree_filter/src/bfs/bfs.cu

@@ -111,13 +111,13 @@ bfs_forward(
     auto adj_vec_len_tensor     = at::zeros({batch_size, vertex_count}, options);
     auto parent_index_tensor    = at::zeros({batch_size, vertex_count}, options);
 
-    int * edge_index      = edge_index_tensor.contiguous().data<int>();
-    int * sorted_index    = sorted_index_tensor.contiguous().data<int>();
-    int * sorted_parent   = sorted_parent_tensor.contiguous().data<int>();
-    int * sorted_child    = sorted_child_tensor.contiguous().data<int>();
-    int * adj_vec         = adj_vec_tensor.contiguous().data<int>();
-    int * adj_vec_len     = adj_vec_len_tensor.contiguous().data<int>();
-    int * parent_index    = parent_index_tensor.contiguous().data<int>();
+    int * edge_index      = edge_index_tensor.contiguous().data_ptr<int>();
+    int * sorted_index    = sorted_index_tensor.contiguous().data_ptr<int>();
+    int * sorted_parent   = sorted_parent_tensor.contiguous().data_ptr<int>();
+    int * sorted_child    = sorted_child_tensor.contiguous().data_ptr<int>();
+    int * adj_vec         = adj_vec_tensor.contiguous().data_ptr<int>();
+    int * adj_vec_len     = adj_vec_len_tensor.contiguous().data_ptr<int>();
+    int * parent_index    = parent_index_tensor.contiguous().data_ptr<int>();
 
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();        
 

furnace/kernels/lib_tree_filter/src/mst/mst.cu

@@ -89,12 +89,12 @@ at::Tensor mst_forward(
     unsigned edge_count = edge_index_tensor.size(1);
 
     auto edge_index_cpu   = edge_index_tensor.cpu();
-    auto edge_weight_cpu  = edge_weight_tensor.cpu(); 
+    auto edge_weight_cpu  = edge_weight_tensor.cpu().to(at::kFloat);
     auto edge_out_cpu     = at::empty({batch_size, vertex_count - 1, 2}, edge_index_cpu.options());
 
-    int * edge_out      = edge_out_cpu.contiguous().data<int>();
-    int * edge_index    = edge_index_cpu.contiguous().data<int>();
-    float * edge_weight = edge_weight_cpu.contiguous().data<float>(); 
+    int * edge_out      = edge_out_cpu.contiguous().data_ptr<int>();
+    int * edge_index    = edge_index_cpu.contiguous().data_ptr<int>();
+    float * edge_weight = edge_weight_cpu.contiguous().data_ptr<float>(); 
 
     // Loop for batch
     std::thread pids[batch_size];

furnace/kernels/lib_tree_filter/src/refine/refine.cu

@@ -14,10 +14,11 @@
 #define CUDA_NUM_THREADS         64
 #define GET_CUDA_CHANNEL(N)      ceil(512.0f / N)
 
+template <typename scalar_t>
 __global__ void root_leaf_prop_kernel(
-        float * in_data, 
-        float * out_data, 
-        float * weight,
+        scalar_t * in_data, 
+        scalar_t * out_data, 
+        scalar_t * weight,
         int * sorted_index, 
         int * sorted_parent_index, 
         int batch_size, 
@@ -53,7 +54,7 @@ __global__ void root_leaf_prop_kernel(
             int par_pos = sorted_index[par];
             for (int k = channel_idx * vertex_count; k < channel_size * vertex_count;
                        k += channel_step * vertex_count){
-                float edge_weight = weight[i];
+                scalar_t edge_weight = weight[i];
                 out_data[cur_pos + k] = in_data[i + k] * (1 - edge_weight * edge_weight) +
                                         out_data[par_pos + k] * edge_weight;
                 __threadfence_block();
@@ -65,10 +66,11 @@ __global__ void root_leaf_prop_kernel(
     }
 }
 
+template <typename scalar_t>
 __global__ void leaf_root_aggr_kernel(
-        float * in_data, 
-        float * out_data, 
-        float * weight,
+        scalar_t * in_data, 
+        scalar_t * out_data, 
+        scalar_t * weight,
         int * sorted_index, 
         int * sorted_child_index, 
         int batch_size, 
@@ -113,7 +115,7 @@ __global__ void leaf_root_aggr_kernel(
             int cur_pos = sorted_index[i];
             for (int k = channel_idx * vertex_count; k < channel_size * vertex_count; 
                     k += channel_step * vertex_count){
-                float aggr_sum;
+                scalar_t aggr_sum;
                 if (in_data != NULL)    
                     aggr_sum = in_data[cur_pos + k];
                 else
@@ -131,13 +133,14 @@ __global__ void leaf_root_aggr_kernel(
     }
 }
 
+template <typename scalar_t>
 __global__ void root_leaf_grad_kernel(
-        float * in_data,
-        float * in_grad,
-        float * out_data,
-        float * out_grad, 
-        float * weight,
-        float * grad,
+        scalar_t * in_data,
+        scalar_t * in_grad,
+        scalar_t * out_data,
+        scalar_t * out_grad, 
+        scalar_t * weight,
+        scalar_t * grad,
         int * sorted_index, 
         int * sorted_parent_index, 
         int batch_size, 
@@ -172,14 +175,14 @@ __global__ void root_leaf_grad_kernel(
         int par_thread  = par % thread_count;
         if ((cur == 0) || (node_per_thread[par_thread] >= par)){
             for (int k = channel_idx; k < channel_size; k += channel_step){
-                float edge_weight   = weight[i];
+                scalar_t edge_weight   = weight[i];
                 int data_offset     = (k % data_channel_size) * vertex_count;
                 int grad_offset     = (k % grad_channel_size) * vertex_count;
                 int out_offset      = k * vertex_count;
 
                 if (cur > 0){
-                    float left  = in_grad[cur + grad_offset] * (out_data[par_pos + data_offset] - edge_weight * in_data[cur + data_offset]);
-                    float right = in_data[cur + data_offset] * (out_grad[par + grad_offset] - edge_weight * in_grad[cur + grad_offset]);
+                    scalar_t left  = in_grad[cur + grad_offset] * (out_data[par_pos + data_offset] - edge_weight * in_data[cur + data_offset]);
+                    scalar_t right = in_data[cur + data_offset] * (out_grad[par + grad_offset] - edge_weight * in_grad[cur + grad_offset]);
 
                     grad[cur + out_offset]      = left + right;
                     out_grad[cur + grad_offset] = in_grad[cur + grad_offset] * (1 - edge_weight * edge_weight) +
@@ -218,31 +221,34 @@ refine_forward(
 
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
-    float * feature_in          = feature_in_tensor.contiguous().data<float>();
-    float * edge_weight         = edge_weight_tensor.contiguous().data<float>();
-    int * sorted_index          = sorted_index_tensor.contiguous().data<int>();
-    int * sorted_parent_index   = sorted_parent_tensor.contiguous().data<int>();
-    int * sorted_child_index    = sorted_child_tensor.contiguous().data<int>();
-    float * feature_aggr        = feature_aggr_tensor.contiguous().data<float>();
-    float * feature_aggr_sum    = feature_aggr_up_tensor.contiguous().data<float>();
-    float * weight_sum          = weight_sum_tensor.contiguous().data<float>();
-    float * weight_aggr_sum     = weight_sum_up_tensor.contiguous().data<float>();
-
-    dim3 feature_block_dims(CUDA_NUM_THREADS, 1, 1), feature_grid_dims(batch_size, channel_size, 1);
-    leaf_root_aggr_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
-            feature_in, feature_aggr_sum, edge_weight, sorted_index, sorted_child_index, batch_size, channel_size, vertex_size, max_adj_per_node);
-    root_leaf_prop_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
-            feature_aggr_sum, feature_aggr, edge_weight, sorted_index, sorted_parent_index, batch_size, channel_size, vertex_size);
-
-    dim3 weight_block_dims(CUDA_NUM_THREADS, 1, 1), weight_grid_dims(batch_size, 1, 1);
-    leaf_root_aggr_kernel <<< weight_grid_dims, weight_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
-            NULL, weight_aggr_sum, edge_weight, sorted_index, sorted_child_index, batch_size, 1, vertex_size, max_adj_per_node);
-    root_leaf_prop_kernel <<< weight_grid_dims, weight_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
-            weight_aggr_sum, weight_sum, edge_weight, sorted_index, sorted_parent_index, batch_size, 1, vertex_size);
-    
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(feature_in_tensor.scalar_type(), "refine_forward", [&] {
+        scalar_t * feature_in          = feature_in_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * edge_weight         = edge_weight_tensor.contiguous().data_ptr<scalar_t>();
+        int * sorted_index             = sorted_index_tensor.contiguous().data_ptr<int>();
+        int * sorted_parent_index      = sorted_parent_tensor.contiguous().data_ptr<int>();
+        int * sorted_child_index       = sorted_child_tensor.contiguous().data_ptr<int>();
+        scalar_t * feature_aggr        = feature_aggr_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * feature_aggr_sum    = feature_aggr_up_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * weight_sum          = weight_sum_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * weight_aggr_sum     = weight_sum_up_tensor.contiguous().data_ptr<scalar_t>();
+
+        dim3 feature_block_dims(CUDA_NUM_THREADS, 1, 1), feature_grid_dims(batch_size, channel_size, 1);
+        leaf_root_aggr_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
+                feature_in, feature_aggr_sum, edge_weight, sorted_index, sorted_child_index, batch_size, channel_size, vertex_size, max_adj_per_node);
+        root_leaf_prop_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
+                feature_aggr_sum, feature_aggr, edge_weight, sorted_index, sorted_parent_index, batch_size, channel_size, vertex_size);
+
+        dim3 weight_block_dims(CUDA_NUM_THREADS, 1, 1), weight_grid_dims(batch_size, 1, 1);
+        leaf_root_aggr_kernel <<< weight_grid_dims, weight_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
+                static_cast<scalar_t *>(NULL), weight_aggr_sum, edge_weight, sorted_index, sorted_child_index, batch_size, 1, vertex_size, max_adj_per_node);
+        root_leaf_prop_kernel <<< weight_grid_dims, weight_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
+                weight_aggr_sum, weight_sum, edge_weight, sorted_index, sorted_parent_index, batch_size, 1, vertex_size);
+    });
+
     auto feature_out_tensor = feature_aggr_tensor / weight_sum_tensor.unsqueeze(1); 
     auto result = std::make_tuple(feature_out_tensor, feature_aggr_tensor, feature_aggr_up_tensor,
-                        weight_sum_tensor, weight_sum_up_tensor);
+            weight_sum_tensor, weight_sum_up_tensor);
+
     return result;
 }
 
@@ -271,28 +277,30 @@ at::Tensor refine_backward_feature(
     const int vertex_size       = feature_in_tensor.size(2);
     const int max_adj_per_node  = sorted_child_tensor.size(2);
 
-    float * feature_in          = feature_in_tensor.contiguous().data<float>();
-    float * edge_weight         = edge_weight_tensor.contiguous().data<float>();
-    int * sorted_index          = sorted_index_tensor.contiguous().data<int>();
-    int * sorted_parent_index   = sorted_parent_tensor.contiguous().data<int>();
-    int * sorted_child_index    = sorted_child_tensor.contiguous().data<int>();
-    float * feature_aggr        = feature_aggr_tensor.contiguous().data<float>();
-    float * feature_aggr_sum    = feature_aggr_up_tensor.contiguous().data<float>();
-    float * weight_sum          = weight_sum_tensor.contiguous().data<float>();
-    float * weight_aggr_sum     = weight_sum_up_tensor.contiguous().data<float>();
-    float * grad_out            = grad_out_tensor.contiguous().data<float>();
-    float * grad_feature        = grad_feature_tensor.contiguous().data<float>();
-
-    float * grad_out_norm           = grad_out_norm_tensor.contiguous().data<float>();
-    float * grad_feature_aggr_sum   = grad_feature_aggr_sum_tensor.contiguous().data<float>();
-
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
-    dim3 feature_block_dims(CUDA_NUM_THREADS, 1, 1), feature_grid_dims(batch_size, channel_size, 1);
-    leaf_root_aggr_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
-            grad_out_norm, grad_feature_aggr_sum, edge_weight, sorted_index, sorted_child_index, batch_size, channel_size, vertex_size, max_adj_per_node);
-    root_leaf_prop_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
-            grad_feature_aggr_sum, grad_feature, edge_weight, sorted_index, sorted_parent_index, batch_size, channel_size, vertex_size);
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(feature_in_tensor.scalar_type(), "refine_backward_feature", [&] {
+        scalar_t * feature_in          = feature_in_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * edge_weight         = edge_weight_tensor.contiguous().data_ptr<scalar_t>();
+        int * sorted_index             = sorted_index_tensor.contiguous().data_ptr<int>();
+        int * sorted_parent_index      = sorted_parent_tensor.contiguous().data_ptr<int>();
+        int * sorted_child_index       = sorted_child_tensor.contiguous().data_ptr<int>();
+        scalar_t * feature_aggr        = feature_aggr_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * feature_aggr_sum    = feature_aggr_up_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * weight_sum          = weight_sum_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * weight_aggr_sum     = weight_sum_up_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * grad_out            = grad_out_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * grad_feature        = grad_feature_tensor.contiguous().data_ptr<scalar_t>();
+
+        scalar_t * grad_out_norm           = grad_out_norm_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * grad_feature_aggr_sum   = grad_feature_aggr_sum_tensor.contiguous().data_ptr<scalar_t>();
+
+        dim3 feature_block_dims(CUDA_NUM_THREADS, 1, 1), feature_grid_dims(batch_size, channel_size, 1);
+        leaf_root_aggr_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
+                grad_out_norm, grad_feature_aggr_sum, edge_weight, sorted_index, sorted_child_index, batch_size, channel_size, vertex_size, max_adj_per_node);
+        root_leaf_prop_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
+                grad_feature_aggr_sum, grad_feature, edge_weight, sorted_index, sorted_parent_index, batch_size, channel_size, vertex_size);
+    });
 
     return grad_feature_tensor;
 }
@@ -318,49 +326,52 @@ at::Tensor refine_backward_weight(
     const int channel_size      = feature_in_tensor.size(1); 
     const int vertex_size       = feature_in_tensor.size(2);
     const int max_adj_per_node  = sorted_child_tensor.size(2);
-
-    float * feature_in          = feature_in_tensor.contiguous().data<float>();
-    float * edge_weight         = edge_weight_tensor.contiguous().data<float>();
-    int * sorted_index          = sorted_index_tensor.contiguous().data<int>();
-    int * sorted_parent_index   = sorted_parent_tensor.contiguous().data<int>();
-    int * sorted_child_index    = sorted_child_tensor.contiguous().data<int>();
-    float * feature_out         = feature_out_tensor.contiguous().data<float>();
-    float * feature_aggr        = feature_aggr_tensor.contiguous().data<float>();
-    float * feature_aggr_sum    = feature_aggr_up_tensor.contiguous().data<float>();
-    float * weight_sum          = weight_sum_tensor.contiguous().data<float>();
-    float * weight_aggr_sum     = weight_sum_up_tensor.contiguous().data<float>();
-    float * grad_out            = grad_out_tensor.contiguous().data<float>();
-    float * grad_weight         = grad_weight_tensor.contiguous().data<float>();
-    
+        
     auto grad_all_channel_tensor        = at::zeros_like(feature_in_tensor, options);
     auto grad_norm_all_channel_tensor   = at::zeros_like(feature_in_tensor, options);
     auto grad_out_norm_aggr_sum_tensor  = at::zeros_like(feature_in_tensor, options);
     auto feature_grad_aggr_sum_tensor   = at::zeros_like(feature_in_tensor, options);
-    
-    float * grad_all_channel            = grad_all_channel_tensor.contiguous().data<float>();
-    float * grad_norm_all_channel       = grad_norm_all_channel_tensor.contiguous().data<float>();
-    float * grad_out_norm_aggr_sum      = grad_out_norm_aggr_sum_tensor.contiguous().data<float>();
-    float * feature_grad_aggr_sum       = feature_grad_aggr_sum_tensor.contiguous().data<float>();
-
-    auto grad_out_norm_tensor = grad_out_tensor / weight_sum_tensor.unsqueeze(1);
-    auto feature_grad_tensor  = grad_out_norm_tensor * feature_out_tensor; 
-    float * grad_out_norm     = grad_out_norm_tensor.contiguous().data<float>();
-    float * feature_grad      = feature_grad_tensor.contiguous().data<float>();
 
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
-    dim3 feature_block_dims(CUDA_NUM_THREADS, 1, 1), feature_grid_dims(batch_size, channel_size, 1);
-    leaf_root_aggr_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
-            grad_out_norm, grad_out_norm_aggr_sum, edge_weight, sorted_index, sorted_child_index, batch_size, channel_size, vertex_size, max_adj_per_node);
-    leaf_root_aggr_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
-            feature_grad, feature_grad_aggr_sum, edge_weight, sorted_index, sorted_child_index, batch_size, channel_size, vertex_size, max_adj_per_node);
-
-    root_leaf_grad_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
-            feature_aggr_sum, grad_out_norm_aggr_sum, feature_aggr, grad_out_norm_aggr_sum, edge_weight, grad_all_channel, 
-            sorted_index, sorted_parent_index, batch_size, channel_size, channel_size, vertex_size);
-    root_leaf_grad_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
-            weight_aggr_sum, feature_grad_aggr_sum, weight_sum, feature_grad_aggr_sum, edge_weight, grad_norm_all_channel, 
-            sorted_index, sorted_parent_index, batch_size, 1, channel_size, vertex_size);
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(feature_in_tensor.scalar_type(), "refine_backward_weight", [&] {
+        scalar_t * feature_in          = feature_in_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * edge_weight         = edge_weight_tensor.contiguous().data_ptr<scalar_t>();
+        int * sorted_index             = sorted_index_tensor.contiguous().data_ptr<int>();
+        int * sorted_parent_index      = sorted_parent_tensor.contiguous().data_ptr<int>();
+        int * sorted_child_index       = sorted_child_tensor.contiguous().data_ptr<int>();
+        scalar_t * feature_out         = feature_out_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * feature_aggr        = feature_aggr_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * feature_aggr_sum    = feature_aggr_up_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * weight_sum          = weight_sum_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * weight_aggr_sum     = weight_sum_up_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * grad_out            = grad_out_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * grad_weight         = grad_weight_tensor.contiguous().data_ptr<scalar_t>();
+        
+        scalar_t * grad_all_channel            = grad_all_channel_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * grad_norm_all_channel       = grad_norm_all_channel_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * grad_out_norm_aggr_sum      = grad_out_norm_aggr_sum_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * feature_grad_aggr_sum       = feature_grad_aggr_sum_tensor.contiguous().data_ptr<scalar_t>();
+
+        auto grad_out_norm_tensor = grad_out_tensor / weight_sum_tensor.unsqueeze(1);
+        auto feature_grad_tensor  = grad_out_norm_tensor * feature_out_tensor; 
+        scalar_t * grad_out_norm     = grad_out_norm_tensor.contiguous().data_ptr<scalar_t>();
+        scalar_t * feature_grad      = feature_grad_tensor.contiguous().data_ptr<scalar_t>();
+
+        dim3 feature_block_dims(CUDA_NUM_THREADS, 1, 1), feature_grid_dims(batch_size, channel_size, 1);
+        leaf_root_aggr_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
+                grad_out_norm, grad_out_norm_aggr_sum, edge_weight, sorted_index, sorted_child_index, batch_size, channel_size, vertex_size, max_adj_per_node);
+        leaf_root_aggr_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
+                feature_grad, feature_grad_aggr_sum, edge_weight, sorted_index, sorted_child_index, batch_size, channel_size, vertex_size, max_adj_per_node);
+
+        root_leaf_grad_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
+                feature_aggr_sum, grad_out_norm_aggr_sum, feature_aggr, grad_out_norm_aggr_sum, edge_weight, grad_all_channel, 
+                sorted_index, sorted_parent_index, batch_size, channel_size, channel_size, vertex_size);
+        root_leaf_grad_kernel <<< feature_grid_dims, feature_block_dims, sizeof(int) * CUDA_NUM_THREADS, stream >>>(
+                weight_aggr_sum, feature_grad_aggr_sum, weight_sum, feature_grad_aggr_sum, edge_weight, grad_norm_all_channel, 
+                sorted_index, sorted_parent_index, batch_size, 1, channel_size, vertex_size);
+
+    });
 
     grad_weight_tensor = (grad_all_channel_tensor - grad_norm_all_channel_tensor).sum(1);