wassname/Flex-Convolution
1
0
Fork 0
mirror of https://github.com/wassname/Flex-Convolution.git synced 2026-09-10 11:40:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add more operation and some refactoring of the code to ease usage.

Browse Source
This commit is contained in:
Fabian Groh
2018-11-06 14:00:16 +01:00
committed by PatWie
parent 32e91a688f
commit 2cbeb7e699
27 changed files with 1264 additions and 452 deletions
Download Patch File Download Diff File Expand all files Collapse all files
user_ops/kernels/cuda_utils.h
+10
View File
@@ -0,0 +1,10 @@
#ifndef LIB_CUDA_UTILS_H_
#define LIB_CUDA_UTILS_H_
template <typename T>
__device__ T* DynamicSharedMemory() {
extern __shared__ __align__(sizeof(T)) unsigned char s_shm[];
return reinterpret_cast<T*>(s_shm);
}
#endif // LIB_CUDA_UTILS_H_
user_ops/kernels/flex_conv_kernel.cc
+2 -1
View File
@@ -70,6 +70,7 @@ struct FlexConvFunctor<CPUDevice, Dtype> {
};
template struct FlexConvFunctor<CPUDevice, float>;
template struct FlexConvFunctor<CPUDevice, double>;
template <typename Dtype>
struct FlexConvGrad<CPUDevice, Dtype> {
@@ -167,7 +168,7 @@ struct FlexConvGrad<CPUDevice, Dtype> {
// template struct FlexConvGrad<CPUDevice, int>;
template struct FlexConvGrad<CPUDevice, float>;
// template struct FlexConvGrad<CPUDevice, double>;
template struct FlexConvGrad<CPUDevice, double>;
} // namespace functor
} // namespace tensorflow
user_ops/kernels/flex_conv_kernel_gpu.cu.cc
+29 -8
View File
@@ -12,7 +12,7 @@ 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.
==============================================================================*/
//Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
#if GOOGLE_CUDA
@@ -20,6 +20,7 @@ limitations under the License.
#include <cub/cub.cuh>
#include "cuda_utils.h"
#include "flex_conv_op.h"
#include "tensorflow/core/util/cuda_kernel_helper.h"
@@ -57,10 +58,10 @@ struct ForwardKernel {
}
__device__ __forceinline__ void operator()() const {
extern __shared__ Dtype s_shm[];
Dtype* s_shm = DynamicSharedMemory<Dtype>();
Dtype* s_theta = (float*)&s_shm[0];
Dtype* s_bias = (float*)&s_shm[Dp * C_Din * C_Dout];
Dtype* s_theta = (Dtype*)&s_shm[0];
Dtype* s_bias = (Dtype*)&s_shm[Dp * C_Din * C_Dout];
// glob ids
int b = blockIdx.z;
@@ -281,7 +282,8 @@ struct BackwardFeatureKernel {
}
__device__ __forceinline__ void operator()() const {
extern __shared__ float s_shm[];
// extern __shared__ float s_shm[];
Dtype* s_shm = DynamicSharedMemory<Dtype>();
int i_n = threadIdx.x;
int i_din = threadIdx.y;
@@ -354,7 +356,10 @@ struct BackwardFeatureKernel {
W += s_bias[i_din * C_Dout + dout_inner];
Dtype value = W * s_topdiff[dout_inner * C_N + i_n];
atomicAdd(
// atomicAdd(
// &d_features_out[b * Din * N + din * N + s_nk[k * C_N + i_n]],
// value);
tensorflow::CudaAtomicAdd(
&d_features_out[b * Din * N + din * N + s_nk[k * C_N + i_n]],
value);
}
@@ -384,6 +389,16 @@ struct BackwardFeatureKernel {
namespace tensorflow {
namespace functor {
template <class Dtype, class NBtype>
struct ForwardKernelType {
typedef FlexConvCuda::ForwardKernel<Dtype, NBtype, 3, 128, 32, 32> type;
};
template <>
struct ForwardKernelType<float, int> {
typedef FlexConvCuda::ForwardKernel<float, int, 3, 128, 32, 64> type;
};
template <typename Dtype>
struct FlexConvFunctor<GPUDevice, Dtype> {
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features,
@@ -399,7 +414,11 @@ struct FlexConvFunctor<GPUDevice, Dtype> {
const int Din = theta.dim_size(2);
const int Dout = theta.dim_size(3);
FlexConvCuda::ForwardKernel<Dtype, NBtype, 3, 128, 32, 64> fwk;
// printf("<f> test: %s\n", __PRETTY_FUNCTION__);
typedef typename ForwardKernelType<Dtype, NBtype>::type FKT;
FKT fwk;
fwk.N = N;
fwk.K = K;
fwk.Din = Din;
@@ -422,6 +441,7 @@ struct FlexConvFunctor<GPUDevice, Dtype> {
};
template struct FlexConvFunctor<GPUDevice, float>;
template struct FlexConvFunctor<GPUDevice, double>;
template <typename Dtype>
struct FlexConvGrad<GPUDevice, Dtype> {
@@ -459,7 +479,7 @@ struct FlexConvGrad<GPUDevice, Dtype> {
const Dtype* topdiff_ptr = reinterpret_cast<const Dtype*>(topdiff.data());
Dtype* grad_features_ptr = reinterpret_cast<float*>(grad_features.data());
Dtype* grad_features_ptr = reinterpret_cast<Dtype*>(grad_features.data());
Dtype* grad_theta_ptr = reinterpret_cast<Dtype*>(grad_theta.data());
Dtype* grad_bias_ptr = reinterpret_cast<Dtype*>(grad_bias.data());
@@ -525,6 +545,7 @@ struct FlexConvGrad<GPUDevice, Dtype> {
};
template struct FlexConvGrad<GPUDevice, float>;
template struct FlexConvGrad<GPUDevice, double>;
} // namespace functor
} // namespace tensorflow
user_ops/kernels/flex_conv_op.cc
+16 -26
View File
@@ -12,7 +12,7 @@ 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.
==============================================================================*/
//Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
#include "flex_conv_op.h"
@@ -91,32 +91,22 @@ class FlexConvGradOp : public OpKernel {
}
};
// Register the CPU kernels.
#define REGISTER_FLEXCONV_OP_CPU(T) \
REGISTER_KERNEL_BUILDER( \
Name("FlexConv").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
FlexConvOp<CPUDevice, T>) \
REGISTER_KERNEL_BUILDER( \
Name("FlexConvGrad").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
FlexConvGradOp<CPUDevice, T>)
#define REGISTER_CUSTOM_OP(NAME, DEVICE, T) \
REGISTER_KERNEL_BUILDER( \
Name(#NAME).Device(DEVICE_##DEVICE).TypeConstraint<T>("T"), \
NAME##Op<DEVICE##Device, T>)
TF_CALL_float(REGISTER_FLEXCONV_OP_CPU);
#undef REGISTER_FLEXCONV_OP_CPU
REGISTER_CUSTOM_OP(FlexConv, CPU, float);
REGISTER_CUSTOM_OP(FlexConvGrad, CPU, float);
REGISTER_CUSTOM_OP(FlexConv, CPU, double);
REGISTER_CUSTOM_OP(FlexConvGrad, CPU, double);
// Register the GPU kernels.
// #ifdef GOOGLE_CUDA
#define REGISTER_FLEXCONV_OP_GPU(T) \
REGISTER_KERNEL_BUILDER( \
Name("FlexConv").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
FlexConvOp<GPUDevice, T>) \
REGISTER_KERNEL_BUILDER( \
Name("FlexConvGrad").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
FlexConvGradOp<GPUDevice, T>)
TF_CALL_float(REGISTER_FLEXCONV_OP_GPU);
#undef REGISTER_FLEXCONV_OP_GPU
// #endif // GOOGLE_CUDA
#ifdef GOOGLE_CUDA
REGISTER_CUSTOM_OP(FlexConv, GPU, float);
REGISTER_CUSTOM_OP(FlexConvGrad, GPU, float);
REGISTER_CUSTOM_OP(FlexConv, GPU, double);
REGISTER_CUSTOM_OP(FlexConvGrad, GPU, double);
#endif // GOOGLE_CUDA
#undef REGISTER_CUSTOM_OP
} // namespace tensorflow
user_ops/kernels/flex_deconv_kernel.cc
+2 -2
View File
@@ -71,6 +71,7 @@ struct FlexDeconvFunctor<CPUDevice, Dtype> {
};
template struct FlexDeconvFunctor<CPUDevice, float>;
template struct FlexDeconvFunctor<CPUDevice, double>;
template <typename Dtype>
struct FlexDeconvGrad<CPUDevice, Dtype> {
@@ -165,9 +166,8 @@ struct FlexDeconvGrad<CPUDevice, Dtype> {
}
};
// template struct FlexDeconvGrad<CPUDevice, int>;
template struct FlexDeconvGrad<CPUDevice, float>;
// template struct FlexDeconvGrad<CPUDevice, double>;
template struct FlexDeconvGrad<CPUDevice, double>;
} // namespace functor
} // namespace tensorflow
user_ops/kernels/flex_deconv_kernel_gpu.cu.cc
+2
View File
@@ -177,6 +177,7 @@ struct FlexDeconvFunctor<GPUDevice, Dtype> {
};
template struct FlexDeconvFunctor<GPUDevice, float>;
template struct FlexDeconvFunctor<GPUDevice, double>;
template <typename Dtype>
struct FlexDeconvGrad<GPUDevice, Dtype> {
@@ -220,6 +221,7 @@ struct FlexDeconvGrad<GPUDevice, Dtype> {
};
template struct FlexDeconvGrad<GPUDevice, float>;
template struct FlexDeconvGrad<GPUDevice, double>;
} // namespace functor
} // namespace tensorflow
user_ops/kernels/flex_deconv_op.cc
+17 -11
View File
@@ -12,7 +12,7 @@ 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.
==============================================================================*/
//Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
#include "flex_deconv_op.h"
@@ -88,16 +88,22 @@ class FlexDeconvGradOp : public OpKernel {
}
};
#define OPNAME(NAME) NAME##Op
#define REGISTER(NAME, Dtype) \
REGISTER_KERNEL_BUILDER( \
Name(#NAME).Device(DEVICE_CPU).TypeConstraint<Dtype>("T"), \
OPNAME(NAME) < CPUDevice, Dtype >); \
REGISTER_KERNEL_BUILDER( \
Name(#NAME).Device(DEVICE_GPU).TypeConstraint<Dtype>("T"), \
OPNAME(NAME) < GPUDevice, Dtype >);
#define REGISTER_CUSTOM_OP(NAME, DEVICE, T) \
REGISTER_KERNEL_BUILDER( \
Name(#NAME).Device(DEVICE_##DEVICE).TypeConstraint<T>("T"), \
NAME##Op<DEVICE##Device, T>)
REGISTER(FlexDeconv, float);
REGISTER(FlexDeconvGrad, float);
REGISTER_CUSTOM_OP(FlexDeconv, CPU, float);
REGISTER_CUSTOM_OP(FlexDeconvGrad, CPU, float);
REGISTER_CUSTOM_OP(FlexDeconv, CPU, double);
REGISTER_CUSTOM_OP(FlexDeconvGrad, CPU, double);
#ifdef GOOGLE_CUDA
REGISTER_CUSTOM_OP(FlexDeconv, GPU, float);
REGISTER_CUSTOM_OP(FlexDeconvGrad, GPU, float);
REGISTER_CUSTOM_OP(FlexDeconv, GPU, double);
REGISTER_CUSTOM_OP(FlexDeconvGrad, GPU, double);
#endif // GOOGLE_CUDA
#undef REGISTER_CUSTOM_OP
} // namespace tensorflow
user_ops/kernels/flex_pool_kernel.cc
+3 -3
View File
@@ -12,7 +12,7 @@ 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.
==============================================================================*/
//Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
#include "flex_pool_op.h"
#include "tensorflow/core/framework/op.h"
@@ -59,6 +59,7 @@ struct FlexPoolFunctor<CPUDevice, Dtype> {
};
template struct FlexPoolFunctor<CPUDevice, float>;
template struct FlexPoolFunctor<CPUDevice, double>;
template <typename Dtype>
struct FlexPoolGrad<CPUDevice, Dtype> {
@@ -94,9 +95,8 @@ struct FlexPoolGrad<CPUDevice, Dtype> {
}
};
// template struct FlexPoolGrad<CPUDevice, int>;
template struct FlexPoolGrad<CPUDevice, float>;
// template struct FlexPoolGrad<CPUDevice, double>;
template struct FlexPoolGrad<CPUDevice, double>;
} // namespace functor
} // namespace tensorflow
user_ops/kernels/flex_pool_kernel_gpu.cu.cc
+12 -18
View File
@@ -12,7 +12,7 @@ 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.
==============================================================================*/
//Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
#if GOOGLE_CUDA
@@ -30,7 +30,7 @@ inline int up2(int len, int th) { return (len - 1) / th + 1; }
template <typename Dtype>
__global__ void forward(const int B, const int N, const int K, const int D,
const Dtype* features, const int* neighborhood,
Dtype* output, int* argmax, float float_min_value) {
Dtype* output, int* argmax, Dtype float_min_value) {
// features: each feature description for each point [B, D, N].
// neighborhood: all K nearest neighbors [B, K, N].
// output: each feature description for each point [B, D, N].
@@ -41,14 +41,14 @@ __global__ void forward(const int B, const int N, const int K, const int D,
d += blockDim.y * gridDim.y) {
for (int n = blockIdx.x * blockDim.x + threadIdx.x; n < N;
n += blockDim.x * gridDim.x) {
float best_value = float_min_value;
Dtype best_value = float_min_value;
int best_id = 0;
const int current_flat = b * D * N + d * N + n;
for (int k_ = 0; k_ < K; ++k_) {
const int other_global_id = neighborhood[b * K * N + k_ * N + n];
const float v = features[b * D * N + d * N + other_global_id];
const Dtype v = features[b * D * N + d * N + other_global_id];
if (best_value < v) {
best_id = other_global_id;
@@ -113,13 +113,9 @@ struct FlexPoolFunctor<GPUDevice, Dtype> {
dim3 grid(up2(N, threads), up2(D, threads), B);
forward<Dtype><<<grid, block>>>(
B, N, K, D,
features_.flat<Dtype>().data(), neighborhood_.flat<int>().data(),
output_->flat<Dtype>().data(), argmax_->flat<int>().data(),
std::numeric_limits<Dtype>::lowest());
B, N, K, D, features_.flat<Dtype>().data(),
neighborhood_.flat<int>().data(), output_->flat<Dtype>().data(),
argmax_->flat<int>().data(), std::numeric_limits<Dtype>::lowest());
if (!ctx->eigen_gpu_device().ok()) {
ctx->SetStatus(
@@ -129,6 +125,7 @@ struct FlexPoolFunctor<GPUDevice, Dtype> {
};
template struct FlexPoolFunctor<GPUDevice, float>;
template struct FlexPoolFunctor<GPUDevice, double>;
template <typename Dtype>
struct FlexPoolGrad<GPUDevice, Dtype> {
@@ -149,13 +146,9 @@ struct FlexPoolGrad<GPUDevice, Dtype> {
grad_features_->NumElements() * sizeof(Dtype));
backward<Dtype><<<grid, block>>>(
B, N, K, D,
features_.flat<Dtype>().data(), neighborhood_.flat<int>().data(),
topdiff_.flat<Dtype>().data(), argmax_.flat<int>().data(),
grad_features_->flat<Dtype>().data());
B, N, K, D, features_.flat<Dtype>().data(),
neighborhood_.flat<int>().data(), topdiff_.flat<Dtype>().data(),
argmax_.flat<int>().data(), grad_features_->flat<Dtype>().data());
if (!ctx->eigen_gpu_device().ok()) {
ctx->SetStatus(tensorflow::errors::Internal("CUDA: FlexPoolGrad Error!"));
@@ -164,6 +157,7 @@ struct FlexPoolGrad<GPUDevice, Dtype> {
};
template struct FlexPoolGrad<GPUDevice, float>;
template struct FlexPoolGrad<GPUDevice, double>;
} // namespace functor
} // namespace tensorflow
user_ops/kernels/flex_pool_op.cc
+14 -26
View File
@@ -12,7 +12,7 @@ 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.
==============================================================================*/
//Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
#include "flex_pool_op.h"
@@ -33,7 +33,6 @@ class FlexPoolOp : public OpKernel {
explicit FlexPoolOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
void Compute(OpKernelContext* ctx) override {
// printf("--> Compute CPU Version <--\n");
const Tensor& features_ = ctx->input(0);
const Tensor& neighborhood_ = ctx->input(1);
@@ -65,7 +64,6 @@ class FlexPoolGradOp : public OpKernel {
explicit FlexPoolGradOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
void Compute(OpKernelContext* ctx) override {
// printf("--> Compute CPU Version <--\n");
const Tensor& features_ = ctx->input(0);
const Tensor& neighborhood_ = ctx->input(1);
const Tensor& topdiff_ = ctx->input(2);
@@ -82,32 +80,22 @@ class FlexPoolGradOp : public OpKernel {
}
};
// Register the CPU kernels.
#define REGISTER_FLEXPOOL_OP_CPU(T) \
REGISTER_KERNEL_BUILDER( \
Name("FlexPool").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
FlexPoolOp<CPUDevice, T>) \
REGISTER_KERNEL_BUILDER( \
Name("FlexPoolGrad").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
FlexPoolGradOp<CPUDevice, T>)
#define REGISTER_CUSTOM_OP(NAME, DEVICE, T) \
REGISTER_KERNEL_BUILDER( \
Name(#NAME).Device(DEVICE_##DEVICE).TypeConstraint<T>("T"), \
NAME##Op<DEVICE##Device, T>)
TF_CALL_float(REGISTER_FLEXPOOL_OP_CPU);
#undef REGISTER_FLEXPOOL_OP_CPU
REGISTER_CUSTOM_OP(FlexPool, CPU, float);
REGISTER_CUSTOM_OP(FlexPoolGrad, CPU, float);
REGISTER_CUSTOM_OP(FlexPool, CPU, double);
REGISTER_CUSTOM_OP(FlexPoolGrad, CPU, double);
// Register the GPU kernels.
#ifdef GOOGLE_CUDA
#define REGISTER_FLEXPOOL_OP_GPU(T) \
REGISTER_KERNEL_BUILDER( \
Name("FlexPool").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
FlexPoolOp<GPUDevice, T>) \
REGISTER_KERNEL_BUILDER( \
Name("FlexPoolGrad").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
FlexPoolGradOp<GPUDevice, T>)
TF_CALL_float(REGISTER_FLEXPOOL_OP_GPU);
#undef REGISTER_FLEXPOOL_OP_GPU
REGISTER_CUSTOM_OP(FlexPool, GPU, float);
REGISTER_CUSTOM_OP(FlexPoolGrad, GPU, float);
REGISTER_CUSTOM_OP(FlexPool, GPU, double);
REGISTER_CUSTOM_OP(FlexPoolGrad, GPU, double);
#endif // GOOGLE_CUDA
#undef REGISTER_CUSTOM_OP
} // namespace tensorflow
user_ops/kernels/knn_bruteforce_kernel.cc
+79
View File
@@ -0,0 +1,79 @@
/* Copyright 2018 ComputerGraphics Tuebingen. All Rights Reserved.
Licensed 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.
==============================================================================*/
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
#include "knn_bruteforce_op.h"
#include "tensorflow/core/framework/op.h"
namespace tensorflow {
namespace functor {
template <typename Dtype, typename NBtype>
struct KnnBruteforceFunctor<CPUDevice, Dtype, NBtype> {
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& positions_,
Tensor* neighborhood_out_, Tensor* distances_,
Tensor* timings) {
// positions [B, Dp, N]
// neighborhood_out_ [B, N, K]
// distances_ [B, N, K]
const auto positions = positions_.tensor<Dtype, 3>();
auto neighborhood_out = neighborhood_out_->tensor<NBtype, 3>();
auto distances_out = distances_->tensor<Dtype, 3>();
const int B = positions_.dim_size(0);
const int Dp = positions_.dim_size(1);
const int N = positions_.dim_size(2);
const int K = neighborhood_out_->dim_size(2);
for (int b = 0; b < B; ++b) {
const Dtype* pc_raw = positions_.flat<Dtype>().data() + b * Dp * N;
Eigen::Map<const Eigen::Matrix<Dtype, Eigen::Dynamic, Eigen::Dynamic,
Eigen::RowMajor>>
pc(pc_raw, Dp, N);
// pc: [Dp, N]
for (int n = 0; n < N; ++n) {
const auto query = pc.col(n);
const auto diffs =
(pc.colwise() - query).array() * (pc.colwise() - query).array();
const auto distances = (diffs.colwise().sum()).array().sqrt();
std::vector<Dtype> vec_dist;
std::vector<Dtype> vec_ids;
for (int i = 0; i < N; ++i) {
vec_dist.push_back(distances(i));
vec_ids.push_back(i);
}
std::sort(std::begin(vec_ids), std::end(vec_ids),
[&](int i1, int i2) { return vec_dist[i1] < vec_dist[i2]; });
for (int k = 0; k < K; ++k) {
neighborhood_out(b, n, k) = vec_ids[k];
distances_out(b, n, k) = vec_dist[vec_ids[k]];
}
}
}
}
};
template struct KnnBruteforceFunctor<CPUDevice, float, int>;
template struct KnnBruteforceFunctor<CPUDevice, double, int>;
} // namespace functor
} // namespace tensorflow
user_ops/kernels/knn_bruteforce_kernel_gpu.cu.cc
+273
View File
@@ -0,0 +1,273 @@
/* Copyright 2018 ComputerGraphics Tuebingen. All Rights Reserved.
Licensed 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.
==============================================================================*/
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
#if GOOGLE_CUDA
#define EIGEN_USE_GPU
#include <vector>
#include "cuda_utils.h"
#include <cuda.h>
#include <helper_cuda.h>
#include <cub/cub.cuh>
#include <curand.h>
#include <curand_kernel.h>
#include <limits>
#include "knn_bruteforce_op.h"
#include "tensorflow/core/util/cuda_kernel_helper.h"
template <typename Dtype, typename NBtype, int C_THREADS = 256, int C_VPT = 2>
struct BlockBFKernel;
template <typename Dtype, typename NBtype, int C_THREADS, int C_VPT>
__global__ void runBlockBFKernel(
const BlockBFKernel<Dtype, NBtype, C_THREADS, C_VPT> kernel) {
kernel();
}
template <typename Dtype, typename NBtype, int C_THREADS, int C_VPT>
struct BlockBFKernel {
void launch(int B) {
dim3 block(C_THREADS, 1);
dim3 grid(N, 1, B);
size_t shm_size = Dp * sizeof(Dtype);
runBlockBFKernel<<<grid, block, shm_size>>>((*this));
}
__device__ __forceinline__ void operator()() const {
// extern __shared__ Dtype s_shm[];
Dtype* s_shm = DynamicSharedMemory<Dtype>();
Dtype* s_point = (Dtype*)&s_shm[0];
if (N > C_THREADS * C_VPT) {
printf(
"<BlockBFKernel> Critical problem!!!!! Not enough resources spend "
"for N points!!! %d < %d \n",
C_THREADS * C_VPT, N);
return;
}
int b = blockIdx.z;
int y = blockIdx.x;
int tid = threadIdx.x;
for (int dpi = tid; dpi < Dp; dpi += blockDim.x) {
s_point[dpi] = d_data[b * N * Dp + dpi * N + y]; // not aligned with
// data!
}
__syncthreads();
Dtype thread_dists[C_VPT]; // keys
NBtype thread_ids[C_VPT]; // values
typedef cub::BlockRadixSort<Dtype, C_THREADS, C_VPT, NBtype> BlockRadixSort;
typedef cub::BlockStore<Dtype, C_THREADS, C_VPT,
cub::BLOCK_STORE_WARP_TRANSPOSE>
BlockStoreDists;
typedef cub::BlockStore<NBtype, C_THREADS, C_VPT,
cub::BLOCK_STORE_WARP_TRANSPOSE>
BlockStoreIds;
// Allocate shared memory
__shared__ union {
typename BlockRadixSort::TempStorage sort;
typename BlockStoreDists::TempStorage store_dists;
typename BlockStoreIds::TempStorage store_ids;
} temp_storage;
for (int vpt_i = 0; vpt_i < C_VPT; ++vpt_i) {
int x = vpt_i * C_THREADS + threadIdx.x;
if (x < N) {
Dtype sum = 0.f;
for (int dpi = 0; dpi < Dp; ++dpi) {
// Dtype val = d_data[b * N * Dp + dpi * N +
//x]
//- d_data[b * N * Dp + dpi * N + y];
Dtype val = d_data[b * N * Dp + dpi * N + x] - s_point[dpi];
sum += val * val;
}
thread_dists[vpt_i] = sqrt(sum);
thread_ids[vpt_i] = x;
} else {
thread_dists[vpt_i] = std::numeric_limits<Dtype>::max();
thread_ids[vpt_i] = -1;
}
}
BlockRadixSort(temp_storage.sort).Sort(thread_dists, thread_ids);
__syncthreads();
BlockStoreIds(temp_storage.store_ids)
.Store(&d_knn_ids[b * N * K + y * K], thread_ids, K);
__syncthreads();
BlockStoreDists(temp_storage.store_dists)
.Store(&d_knn_dists[b * N * K + y * K], thread_dists, K);
}
const Dtype* d_data;
Dtype* d_knn_dists;
NBtype* d_knn_ids;
int N;
int Dp;
int K;
};
template <typename Dtype, typename NBtype>
struct BlockBFKernelAttributesSetter {
template <typename T>
void setAttributes(T& kernel) {
kernel.d_data = d_data;
kernel.d_knn_dists = d_knn_dists;
kernel.d_knn_ids = d_knn_ids;
kernel.N = N;
kernel.Dp = Dp;
kernel.K = K;
}
const Dtype* d_data;
Dtype* d_knn_dists;
NBtype* d_knn_ids;
int N;
int Dp;
int K;
};
namespace tensorflow {
namespace functor {
template <typename Dtype, typename NBtype>
struct KnnBruteforceFunctor<GPUDevice, Dtype, NBtype> {
void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& positions,
Tensor* neighborhood_out, Tensor* distances,
Tensor* timings) {
// printf("GPU: KNNBF! \n");
// printf("return_timings: %d \n",return_timings);
const int B = positions.dim_size(0);
const int D = positions.dim_size(1);
const int N = positions.dim_size(2);
const int K = neighborhood_out->dim_size(2);
// printf("B: %d | N: %d | K: %d \n", B, N, K);
BlockBFKernelAttributesSetter<Dtype, NBtype> attr;
attr.d_data = positions.flat<Dtype>().data();
attr.d_knn_ids = neighborhood_out->flat<NBtype>().data();
attr.d_knn_dists = distances->flat<Dtype>().data();
attr.N = N;
attr.Dp = D;
attr.K = K;
cudaEvent_t start, stop;
if (return_timings) {
cudaEventCreate(&start);
cudaEventCreate(&stop);
cudaEventRecord(start, 0);
}
// TODO: beautify
if (N <= 32) {
BlockBFKernel<Dtype, NBtype, 32, 1> knn;
attr.setAttributes(knn);
knn.launch(B);
} else if (N <= 64) {
BlockBFKernel<Dtype, NBtype, 64, 1> knn;
attr.setAttributes(knn);
knn.launch(B);
} else if (N <= 128) {
BlockBFKernel<Dtype, NBtype, 128, 1> knn;
attr.setAttributes(knn);
knn.launch(B);
} else if (N <= 256) {
BlockBFKernel<Dtype, NBtype, 128, 2> knn;
attr.setAttributes(knn);
knn.launch(B);
} else if (N <= 512) {
BlockBFKernel<Dtype, NBtype, 128, 4> knn;
attr.setAttributes(knn);
knn.launch(B);
} else if (N <= 1024) {
BlockBFKernel<Dtype, NBtype, 256, 4> knn;
attr.setAttributes(knn);
knn.launch(B);
} else if (N <= 2048) {
BlockBFKernel<Dtype, NBtype, 256, 8> knn;
attr.setAttributes(knn);
knn.launch(B);
} else if (N <= 4096) {
BlockBFKernel<Dtype, NBtype, 512, 8> knn;
attr.setAttributes(knn);
knn.launch(B);
} else if (N <= 1024 * 8) {
BlockBFKernel<Dtype, NBtype, 1024, 8> knn;
attr.setAttributes(knn);
knn.launch(B);
} else {
printf(
"point sets greater then 8k are note yet supported!! Change to "
"knn_graph operation!! \n");
}
if (return_timings) {
float time;
cudaEventRecord(stop, 0);
cudaEventSynchronize(stop);
cudaEventElapsedTime(&time, start, stop);
// printf("complete elapsed time for KNNBf: %f ms
//\n", time);
Dtype time2 = time;
cudaMemcpy(timings->flat<Dtype>().data(), &time2, sizeof(Dtype),
cudaMemcpyHostToDevice);
cudaEventDestroy(start);
cudaEventDestroy(stop);
}
cudaDeviceSynchronize();
getLastCudaError("KNNBF execution failed");
checkCudaErrors(cudaDeviceSynchronize());
}
bool return_timings;
};
template struct KnnBruteforceFunctor<GPUDevice, float, int>;
// // too much shared memory
// template struct KnnBruteforceFunctor<GPUDevice, double, int>;
} // namespace functor
} // namespace tensorflow
#endif // GOOGLE_CUDA
user_ops/kernels/knn_bruteforce_op.cc
+81
View File
@@ -0,0 +1,81 @@
/* Copyright 2018 ComputerGraphics Tuebingen. All Rights Reserved.
Licensed 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.
==============================================================================*/
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
#include "knn_bruteforce_op.h"
#include <stdio.h>
#include <type_traits>
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
namespace tensorflow {
// Forward-Pass (CPU, GPU)
// --------------------------------------------------
template <typename Device, typename Dtype>
class KnnBruteforceOp : public OpKernel {
public:
explicit KnnBruteforceOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("K", &K));
OP_REQUIRES_OK(ctx, ctx->GetAttr("return_timings", &return_timings));
}
void Compute(OpKernelContext* ctx) override {
const Tensor& positions = ctx->input(0);
const Tensor& neighborhood_in = ctx->input(1);
const int B = positions.shape().dim_size(0);
const int N = positions.shape().dim_size(2);
Tensor* neighborhood_out = nullptr;
Tensor* distances = nullptr;
Tensor* timings = nullptr;
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({B, N, K}),
&neighborhood_out));
OP_REQUIRES_OK(ctx,
ctx->allocate_output(1, TensorShape({B, N, K}), &distances));
OP_REQUIRES_OK(ctx, ctx->allocate_output(2, TensorShape({1}), &timings));
::tensorflow::functor::KnnBruteforceFunctor<Device, Dtype, int> knnBFF;
knnBFF.return_timings = return_timings;
knnBFF(ctx, positions, neighborhood_out, distances, timings);
}
private:
TF_DISALLOW_COPY_AND_ASSIGN(KnnBruteforceOp);
int K;
bool return_timings;
// int subBatch;
};
#define REGISTER_CUSTOM_OP(NAME, DEVICE, T) \
REGISTER_KERNEL_BUILDER( \
Name(#NAME).Device(DEVICE_##DEVICE).TypeConstraint<T>("T"), \
NAME##Op<DEVICE##Device, T>)
REGISTER_CUSTOM_OP(KnnBruteforce, CPU, float);
#ifdef GOOGLE_CUDA
REGISTER_CUSTOM_OP(KnnBruteforce, GPU, float);
#endif // GOOGLE_CUDA
#undef REGISTER_CUSTOM_OP
} // namespace tensorflow
user_ops/kernels/knn_bruteforce_op.h
+44
View File
@@ -0,0 +1,44 @@
/* Copyright 2018 ComputerGraphics Tuebingen. All Rights Reserved.
Licensed 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.
==============================================================================*/
// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
#ifndef LIB_KNN_BF_OP_H_
#define LIB_KNN_BF_OP_H_
#include "tensorflow/core/framework/op_kernel.h"
namespace tensorflow {
class OpKernelContext;
class Tensor;
using CPUDevice = Eigen::ThreadPoolDevice;
using GPUDevice = Eigen::GpuDevice;
} // namespace tensorflow
namespace tensorflow {
namespace functor {
template <typename Device, typename Dtype, typename NBtype>
struct KnnBruteforceFunctor {
void operator()(::tensorflow::OpKernelContext *ctx, const Tensor &positions,
Tensor *neighborhood_out, Tensor *distances, Tensor *timings);
bool return_timings;
};
} // namespace functor
} // namespace tensorflow
#endif // LIB_KNN_BF_OP_H_