kdtree_cuda_3d_index.h

Go to the documentation of this file.

/***********************************************************************
 * Software License Agreement (BSD License)
 *
 * Copyright 2008-2009  Marius Muja (mariusm@cs.ubc.ca). All rights reserved.
 * Copyright 2008-2009  David G. Lowe (lowe@cs.ubc.ca). All rights reserved.
 * Copyright 2011       Andreas Muetzel (amuetzel@uni-koblenz.de). All rights reserved.
 *
 * THE BSD LICENSE
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *************************************************************************/
 
#ifndef FLANN_KDTREE_CUDA_3D_INDEX_H_
#define FLANN_KDTREE_CUDA_3D_INDEX_H_
 
#include <algorithm>
#include <map>
#include <cassert>
#include <cstring>
#include "FLANN/general.h"
#include "FLANN/algorithms/nn_index.h"
#include "FLANN/util/matrix.h"
#include "FLANN/util/result_set.h"
#include "FLANN/util/heap.h"
#include "FLANN/util/allocator.h"
#include "FLANN/util/random.h"
#include "FLANN/util/saving.h"
#include "FLANN/util/params.h"
 
namespace flann
{
 
struct KDTreeCuda3dIndexParams : public IndexParams
{
    KDTreeCuda3dIndexParams( int leaf_max_size = 64 )
    {
        (*this)["algorithm"] = FLANN_INDEX_KDTREE_CUDA;
        (*this)["leaf_max_size"] = leaf_max_size;
        (*this)["dim"] = 3;
    }
};
 
template <typename Distance>
class KDTreeCuda3dIndex : public NNIndex<Distance>
{
public:
    typedef typename Distance::ElementType ElementType;
    typedef typename Distance::ResultType DistanceType;
 
    typedef NNIndex<Distance> BaseClass;
 
    int visited_leafs;
 
    typedef bool needs_kdtree_distance;
 
    KDTreeCuda3dIndex(const Matrix<ElementType>& inputData, const IndexParams& params = KDTreeCuda3dIndexParams(),
                      Distance d = Distance() ) : BaseClass(params,d), dataset_(inputData), leaf_count_(0), visited_leafs(0), node_count_(0), current_node_count_(0)
    {
        size_ = dataset_.rows;
        dim_ = dataset_.cols;
 
        int dim_param = get_param(params,"dim",-1);
        if (dim_param>0) dim_ = dim_param;
        leaf_max_size_ = get_param(params,"leaf_max_size",10);
        assert( dim_ == 3 );
        gpu_helper_=0;
    }
 
    KDTreeCuda3dIndex(const KDTreeCuda3dIndex& other);
    KDTreeCuda3dIndex operator=(KDTreeCuda3dIndex other);
 
    ~KDTreeCuda3dIndex()
    {
        delete[] data_.ptr();
        clearGpuBuffers();
    }
 
    BaseClass* clone() const
    {
        throw FLANNException("KDTreeCuda3dIndex cloning is not implemented");
    }
 
    void buildIndex()
    {
        // Create a permutable array of indices to the input vectors.
        vind_.resize(size_);
        for (size_t i = 0; i < size_; i++) {
            vind_[i] = i;
        }
 
        leaf_count_=0;
        node_count_=0;
        //         computeBoundingBox(root_bbox_);
        //                              tree_.reserve(log2((double)size_/leaf_max_size_));
        //         divideTree(0, size_, root_bbox_,-1 );   // construct the tree
 
        delete[] data_.ptr();
 
        uploadTreeToGpu();
    }
 
    flann_algorithm_t getType() const
    {
        return FLANN_INDEX_KDTREE_SINGLE;
    }
 
 
    void removePoint(size_t index)
    {
        throw FLANNException( "removePoint not implemented for this index type!" );
    }
 
    ElementType* getPoint(size_t id)
    {
        return dataset_[id];
    }
 
    void saveIndex(FILE* stream)
    {
        throw FLANNException( "Index saving not implemented!" );
 
    }
 
 
    void loadIndex(FILE* stream)
    {
        throw FLANNException( "Index loading not implemented!" );
    }
 
    size_t veclen() const
    {
        return dim_;
    }
 
    int usedMemory() const
    {
        //         return tree_.size()*sizeof(Node)+dataset_.rows*sizeof(int);  // pool memory and vind array memory
        return 0;
    }
 
 
    int knnSearch(const Matrix<ElementType>& queries, Matrix<int>& indices, Matrix<DistanceType>& dists, size_t knn, const SearchParams& params) const
    {
        knnSearchGpu(queries,indices, dists, knn, params);
        return knn*queries.rows; // hack...
    }
 
    int knnSearch(const Matrix<ElementType>& queries,
                          std::vector< std::vector<int> >& indices,
                          std::vector<std::vector<DistanceType> >& dists,
                          size_t knn,
                          const SearchParams& params) const
    {
        knnSearchGpu(queries,indices, dists, knn, params);
        return knn*queries.rows; // hack...
    }
 
    void knnSearchGpu(const Matrix<ElementType>& queries, Matrix<int>& indices, Matrix<DistanceType>& dists, size_t knn, const SearchParams& params) const;
 
    int knnSearchGpu(const Matrix<ElementType>& queries,
                     std::vector< std::vector<int> >& indices,
                     std::vector<std::vector<DistanceType> >& dists,
                     size_t knn,
                     const SearchParams& params) const
    {
        flann::Matrix<int> ind( new int[knn*queries.rows], queries.rows,knn);
        flann::Matrix<DistanceType> dist( new DistanceType[knn*queries.rows], queries.rows,knn);
        knnSearchGpu(queries,ind,dist,knn,params);
        for( size_t i = 0; i<queries.rows; i++ ) {
            indices[i].resize(knn);
            dists[i].resize(knn);
            for( size_t j=0; j<knn; j++ ) {
                indices[i][j]=ind[i][j];
                dists[i][j]=dist[i][j];
            }
        }
        delete [] ind.ptr();
        delete [] dist.ptr();
        return knn*queries.rows; // hack...
    }
 
    int radiusSearch(const Matrix<ElementType>& queries, Matrix<int>& indices, Matrix<DistanceType>& dists,
                             float radius, const SearchParams& params) const
    {
        return radiusSearchGpu(queries,indices, dists, radius, params);
    }
 
    int radiusSearch(const Matrix<ElementType>& queries, std::vector< std::vector<int> >& indices,
                             std::vector<std::vector<DistanceType> >& dists, float radius, const SearchParams& params) const
    {
        return radiusSearchGpu(queries,indices, dists, radius, params);
    }
 
    int radiusSearchGpu(const Matrix<ElementType>& queries, Matrix<int>& indices, Matrix<DistanceType>& dists,
                        float radius, const SearchParams& params) const;
 
    int radiusSearchGpu(const Matrix<ElementType>& queries, std::vector< std::vector<int> >& indices,
                        std::vector<std::vector<DistanceType> >& dists, float radius, const SearchParams& params) const;
 
    void findNeighbors(ResultSet<DistanceType>& result, const ElementType* vec, const SearchParams& searchParams) const
    {
    }
 
protected:
    void buildIndexImpl()
    {
        /* nothing to do here */
    }
 
    void freeIndex()
    {
        /* nothing to do here */
    }
 
private:
 
    void uploadTreeToGpu( );
 
    void clearGpuBuffers( );
 
 
 
 
private:
 
    struct GpuHelper;
 
    GpuHelper* gpu_helper_;
 
    const Matrix<ElementType> dataset_;
 
    int leaf_max_size_;
 
    int leaf_count_;
    int node_count_;
    int current_node_count_;
 
 
    std::vector<int> vind_;
 
    Matrix<ElementType> data_;
 
    size_t dim_;
 
    USING_BASECLASS_SYMBOLS
};   // class KDTreeCuda3dIndex
 
 
}
 
#endif //FLANN_KDTREE_SINGLE_INDEX_H_