Method Run

tensorflow/core/kernels/random_op_gpu.h:133–159

// distribution. Each output takes a fixed number of samples.
template <class Distribution>
PHILOX_DEVICE_INLINE void FillPhiloxRandomKernel<Distribution, false>::Run(
    random::PhiloxRandom gen, T* data, int64 size, Distribution dist) {
  const int kGroupSize = Distribution::kResultElementCount;

  const int32 thread_id = blockIdx.x * blockDim.x + threadIdx.x;
  const int32 total_thread_count = gridDim.x * blockDim.x;
  int64 offset = thread_id * kGroupSize;
  gen.Skip(thread_id);

  const SampleCopier<T, kGroupSize> copier;
  while (offset + kGroupSize <= size) {
    const typename Distribution::ResultType samples = dist(&gen);
    copier(&data[offset], samples);

    offset += total_thread_count * kGroupSize;
    gen.Skip(total_thread_count - 1);
  }

  typename Distribution::ResultType samples = dist(&gen);
  for (int i = 0; i < kGroupSize; ++i) {
    if (offset >= size) {
      return;
    }
    data[offset] = samples[i];
    ++offset;
  }
}

// A cuda kernel to fill the data with random numbers from the specified
// distribution. Each output takes a variable number of samples.