Pushing changes to GitHub Pages.

NVIDIA-Merlin · Aug 20, 2024 · f267e0a · f267e0a
1 parent b0d1a78
commit f267e0a
Show file tree

Hide file tree

Showing 6 changed files with 63 additions and 11 deletions.
diff --git a/review/pr-200/_sources/api/program_listing_file_merlin_hashtable.cuh.rst.txt b/review/pr-200/_sources/api/program_listing_file_merlin_hashtable.cuh.rst.txt
@@ -363,6 +363,7 @@ Program Listing for File merlin_hashtable.cuh
        CUDA_CHECK(cudaGetDeviceProperties(&deviceProp, options_.device_id));
        shared_mem_size_ = deviceProp.sharedMemPerBlock;
        sm_cnt_ = deviceProp.multiProcessorCount;
+       max_threads_per_block_ = deviceProp.maxThreadsPerBlock;
        create_table<key_type, value_type, score_type>(
            &table_, allocator_, options_.dim, options_.init_capacity,
            options_.max_capacity, options_.max_hbm_for_vectors,
@@ -1627,10 +1628,10 @@ Program Listing for File merlin_hashtable.cuh
                               offset % TILE_SIZE == 0 && n % TILE_SIZE == 0;
    
        if (match_fast_cond) {
-         int grid_size = std::min(sm_cnt_, static_cast<int>(SAFE_GET_GRID_SIZE(
-                                               n, options_.block_size)));
-         const int TILE_SIZE = 8;
-   
+         int grid_size = std::min(
+             sm_cnt_ * max_threads_per_block_ / options_.block_size,
+             static_cast<int>(SAFE_GET_GRID_SIZE(n, options_.block_size)));
+         const int TILE_SIZE = 32;
          dump_kernel_v2<key_type, value_type, score_type, PredFunctor, TILE_SIZE>
              <<<grid_size, options_.block_size, 0, stream>>>(
                  d_table_, table_->buckets, pattern, threshold, keys, values,
@@ -1681,6 +1682,22 @@ Program Listing for File merlin_hashtable.cuh
        return h_size;
      }
    
+     template <template <typename, typename> class PredFunctor>
+     void size_if(const key_type& pattern, const score_type& threshold,
+                  size_type* d_counter, cudaStream_t stream = 0) const {
+       read_shared_lock lock(mutex_, stream);
+       CUDA_CHECK(cudaMemsetAsync(d_counter, 0, sizeof(size_type), stream));
+   
+       size_t grid_size = SAFE_GET_GRID_SIZE(capacity(), options_.block_size);
+       grid_size = std::min(grid_size,
+                            static_cast<size_t>(sm_cnt_ * max_threads_per_block_ /
+                                                options_.block_size));
+       size_if_kernel<key_type, value_type, score_type, PredFunctor>
+           <<<grid_size, options_.block_size, 0, stream>>>(
+               d_table_, table_->buckets, pattern, threshold, d_counter);
+       CudaCheckError();
+     }
+   
      size_type capacity() const { return table_->capacity; }
    
      void reserve(const size_type new_capacity, cudaStream_t stream = 0) {
@@ -1958,6 +1975,7 @@ Program Listing for File merlin_hashtable.cuh
      TableCore* d_table_ = nullptr;
      size_t shared_mem_size_ = 0;
      int sm_cnt_ = 0;
+     int max_threads_per_block_ = 0;
      std::atomic_bool reach_max_capacity_{false};
      bool initialized_ = false;
      mutable group_shared_mutex mutex_;

diff --git a/review/pr-200/api/classnv_1_1merlin_1_1HashTable.html b/review/pr-200/api/classnv_1_1merlin_1_1HashTable.html
@@ -623,6 +623,20 @@ <h2>Class Documentation<a class="headerlink" href="#class-documentation" title="
 </dl>
 </dd></dl>
 
+<dl class="cpp function">
+<dt id="_CPPv4II00E0ENK2nv6merlin9HashTable7size_ifEvRK8key_typeRK10score_typeP9size_type12cudaStream_t">
+<span id="_CPPv3II00E0ENK2nv6merlin9HashTable7size_ifERK8key_typeRK10score_typeP9size_type12cudaStream_t"></span><span id="_CPPv2II00E0ENK2nv6merlin9HashTable7size_ifERK8key_typeRK10score_typeP9size_type12cudaStream_t"></span><span class="pre">template&lt;</span><span class="pre">template&lt;</span><span class="pre">typename</span><span class="pre">,</span> <span class="pre">typename</span><span class="pre">&gt;</span> <span class="pre">class</span> <code class="sig-name descname"><span class="pre">PredFunctor</span></code><span class="pre">&gt;</span><br /><span class="target" id="classnv_1_1merlin_1_1HashTable_1a3a5a38ba3cd6b68b707892ea28274ddf"></span><em class="property"><span class="pre">inline</span></em> <span class="pre">void</span> <code class="sig-name descname"><span class="pre">size_if</span></code><span class="sig-paren">(</span><em class="property"><span class="pre">const</span></em> <a class="reference internal" href="#_CPPv4N2nv6merlin9HashTable8key_typeE" title="nv::merlin::HashTable::key_type"><span class="pre">key_type</span></a> <span class="pre">&amp;</span><em><span class="pre">pattern</span></em>, <em class="property"><span class="pre">const</span></em> <a class="reference internal" href="#_CPPv4N2nv6merlin9HashTable10score_typeE" title="nv::merlin::HashTable::score_type"><span class="pre">score_type</span></a> <span class="pre">&amp;</span><em><span class="pre">threshold</span></em>, <a class="reference internal" href="#_CPPv4N2nv6merlin9HashTable9size_typeE" title="nv::merlin::HashTable::size_type"><span class="pre">size_type</span></a> <span class="pre">*</span><em><span class="pre">d_counter</span></em>, <span class="pre">cudaStream_t</span> <em><span class="pre">stream</span></em> <span class="pre">=</span> <span class="pre">0</span><span class="sig-paren">)</span> <em class="property"><span class="pre">const</span></em><a class="headerlink" href="#_CPPv4II00E0ENK2nv6merlin9HashTable7size_ifEvRK8key_typeRK10score_typeP9size_type12cudaStream_t" title="Permalink to this definition"></a><br /></dt>
+<dd><p>Returns the number of keys if meet PredFunctor. </p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><p><strong>stream</strong> – The CUDA stream that is used to execute the operation. </p>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p>The table size match condiction of PredFunctor. </p>
+</dd>
+</dl>
+</dd></dl>
+
 <dl class="cpp function">
 <dt id="_CPPv4NK2nv6merlin9HashTable8capacityEv">
 <span id="_CPPv3NK2nv6merlin9HashTable8capacityEv"></span><span id="_CPPv2NK2nv6merlin9HashTable8capacityEv"></span><span id="nv::merlin::HashTable::capacityC"></span><span class="target" id="classnv_1_1merlin_1_1HashTable_1a083051d75a459ec2ee6a94257ff71457"></span><em class="property"><span class="pre">inline</span></em> <a class="reference internal" href="#_CPPv4N2nv6merlin9HashTable9size_typeE" title="nv::merlin::HashTable::size_type"><span class="pre">size_type</span></a> <code class="sig-name descname"><span class="pre">capacity</span></code><span class="sig-paren">(</span><span class="sig-paren">)</span> <em class="property"><span class="pre">const</span></em><a class="headerlink" href="#_CPPv4NK2nv6merlin9HashTable8capacityEv" title="Permalink to this definition"></a><br /></dt>