submission 68656

  #!POPCORN leaderboard histogram_v2
   
- import triton
- import triton.language as tl
+ from torch.utils.cpp_extension import load_inline
  from task import input_t, output_t
+ import torch
   
   
- @triton.jit
- def kernel(
-     data_ptr,  # (size,)
-     output_ptr,  # (256,)
-     size,
-     BLOCK_SIZE: tl.constexpr,
-     NUM_BINS: tl.constexpr = 256,
- ):
-     pid = tl.program_id(0)
-     num_pids = tl.num_programs(0)
+ CUDA_SRC = r"""
+ constexpr int WARP_SIZE = 32;
+ constexpr int NUM_BINS = 256;
+ constexpr int NUM_WARPS = 8;
+ constexpr int TB_SIZE = NUM_WARPS * WARP_SIZE;
   
-     acc = tl.zeros((NUM_BINS,), dtype=tl.int32)
+ __device__ __host__
+ constexpr int cdiv(int a, int b) { return (a + b - 1) / b; }
   
-     num_iters = tl.cdiv(size, BLOCK_SIZE * num_pids)
-     for iter_id in range(num_iters):
-         offs = iter_id * (num_pids * BLOCK_SIZE) + (pid * BLOCK_SIZE) + tl.arange(0, BLOCK_SIZE)
-         mask = offs < size
-         data = tl.load(data_ptr + offs, mask, other=0).to(tl.int32)  # tl.histogram() doesn't work with uint8
-         acc += tl.histogram(data, NUM_BINS)  # old triton doesn't have mask for histogram
+ __align__(16)
+ struct u8x16 { uint8_t x[16]; };
   
-     # NOTE: output_ptr is i64 type
-     tl.atomic_add(output_ptr + tl.arange(0, NUM_BINS), acc)
+ __global__
+ void kernel(
+   const uint8_t *data_ptr,    // (size,)
+         int64_t *output_ptr,  // (256,)
+         int64_t size) {
   
-     # compensation since we use 0 for masked elements
-     if pid == 0:
-         compensate = size - num_iters * BLOCK_SIZE * num_pids
-         tl.atomic_add(output_ptr, compensate)
+   const int tid = threadIdx.x;
+   const int bid = blockIdx.x;
+   const int num_blocks = gridDim.x;
+   const int warp_id = tid / WARP_SIZE;
+   const int lane_id = tid % WARP_SIZE;
   
+   __shared__ int smem_hist[NUM_WARPS * NUM_BINS];
   
+   // init
+   for (int iter_id = 0; iter_id < (NUM_WARPS * NUM_BINS / TB_SIZE); iter_id++)
+     smem_hist[iter_id * TB_SIZE + tid] = 0;
+   __syncthreads();
+  
+   // each thread reads 16 elems
+   const u8x16 *data_u8x16_ptr = reinterpret_cast<const u8x16 *>(data_ptr);
+   data_u8x16_ptr += bid * TB_SIZE + tid;
+  
+   const int wave_size = (16 * TB_SIZE * num_blocks);
+   const int num_iters = size / wave_size;
+   for (int iter_id = 0; iter_id < num_iters; iter_id ++) {
+     const u8x16 x = data_u8x16_ptr[0];
+     data_u8x16_ptr += num_blocks * TB_SIZE;
+  
+     for (int elem_id = 0; elem_id < 16; elem_id++) {
+       const int val = x.x[elem_id];  // cast u8->i32
+       atomicAdd(smem_hist + (warp_id * NUM_BINS + val), 1);
+     }
+   }
+  
+   // each thread reads 1 elem
+   const int start = num_iters * wave_size + bid * TB_SIZE + tid;
+   for (int i = start; i < size; i += num_blocks * TB_SIZE) {
+     const int val = data_ptr[i];
+     atomicAdd(smem_hist + (warp_id * NUM_BINS + val), 1);
+   }
+  
+   __syncthreads();
+  
+   // combine histogram across warps
+   static_assert(NUM_BINS % TB_SIZE == 0);
+   for (int iter_id = 0; iter_id < NUM_BINS / TB_SIZE; iter_id++) {
+     const int bin_id = iter_id * TB_SIZE + tid;
+     int count = smem_hist[bin_id];  // from 1st sub-histogram
+  
+     for (int sub_id = 1; sub_id < NUM_WARPS; sub_id++)
+       count += smem_hist[sub_id * NUM_BINS + bin_id];
+  
+     // total count shouldn't exceed int32...
+     atomicAdd(reinterpret_cast<int *>(output_ptr + bin_id), count);
+   }
+ }
+  
+ void launch(const at::Tensor& data, at::Tensor& output) {
+   output.zero_();
+  
+   const auto data_ptr = data.data_ptr<uint8_t>();
+   auto output_ptr = output.data_ptr<int64_t>();
+   const int64_t size = data.size(0);
+  
+   const int num_blocks = 264;
+   kernel<<<num_blocks, TB_SIZE>>>(data_ptr, output_ptr, size);
+ }
+  
+ TORCH_LIBRARY(my_module, m) {
+   m.def("launch(Tensor data, Tensor(a!) output) -> ()");
+   m.impl("launch", &launch);
+ }
+ """
+  
+ load_inline(
+     "histogram_v0",
+     cpp_sources="",
+     cuda_sources=CUDA_SRC,
+     verbose=True,
+     is_python_module=False,
+ )
+  
+  
  def custom_kernel(data: input_t) -> output_t:
      data, output = data
+     torch.ops.my_module.launch(data, output)
   
-     BLOCK_SIZE = 8192
-     num_blocks = 264
-     output.zero_()
-     kernel[(num_blocks,)](data, output, data.shape[0], BLOCK_SIZE)
+     # print(output)
+     # print(data.bincount(minlength=256))
   
      return output