code: def _threefry(
      irb,
      key_buf,
      key_offset,
      counter_buf,
      counter_offset,
      out_buf,
      out_offset,
      out_shape,
  ):
      """IRBuilder code for running Threefry
  
      Parameters
      ----------
      irb: IRBuilder
          IRBuilder that this code will be generated for.
  
      key_buf: BufferVar
          Buffer to read the key from.
  
      key_offset: number
          Threefry will write to :code:`key_buf[key_offset:key_offset+4]`
  
      counter_buf: BufferVar
          Buffer to read the counter from.
  
      counter_offset: number
          Threefry will write to :code:`counter_buf[counter_offset:counter_offset+4]`
  
      out_buf: BufferVar
          Buffer to read the counter from.
  
      out_offset: number
          Threefry will write to :code:`out_buf[out_offset:out_offset+4*product(out_shape)]`
  
      out_shape: number
          Determines the number of output states to generate. :code:`state[i]` will correspond to
          counter+i.
      """
      nrounds = 20
      nwords = 4
      iwidth = 64
      assert nrounds % 4 == 0
      assert nwords in [4, 8, 16]
  
      # The paper has constants for 32 bit threefry, but we keep the implementation simple by only
      # using 64-bit words.
      assert (
          key_buf.dtype == "uint64"
      ), "threefry only supports 64-bit keys"
      assert (
          key_buf.dtype == counter_buf.dtype
      ), "threefry key and counter must be the same dtype"
  
      def mix(a, b, rotation):
          x = a + b  # wrapping
          y = x ^ ((b << rotation) | (b >> (iwidth - rotation)))
          return [x, y]
  
      # temporary buffer for holding the results of _PERMUTATIONS
      tmp = irb.allocate(
          out_buf.dtype, out_shape * nwords, name="tmp", scope="global"
      )
      tmp_offset = 0
  
      # Initialize entire key. It is composed of the original key with one
      # element appended. The appended element is the xor of all key words plus a
      # constant.
      full_key = irb.allocate(
          "uint64", nwords + 1, name="full_key", scope="global"
      )
      for i in range(nwords):
          full_key[i] = key_buf[key_offset + i]
      # initial key constant, full_key[nwords] is equivalent to k_{N_W} in the Skein paper.
      full_key[nwords] = tvm.tir.const(
          0x1BD11BDAA9FC1A22, dtype="uint64"
      )
      for i in range(nwords):
          full_key[nwords] ^= key_buf[key_offset + i]
  
      with irb.for_range(0, out_shape, dtype="uint64", name="i") as i:
          for j in range(nwords):
              out_buf[out_offset + i * nwords + j] = (
                  counter_buf[counter_offset + j] + i
              )
  
      def key_schedule(s, i):
          # Threefry uses no tweak, so the key schedule is simple
          if i == nwords - 1:
              return full_key[(s + i) % (nwords + 1)] + tvm.tir.const(
                  s, dtype="uint64"
              )
          return full_key[(s + i) % (nwords + 1)]
  
      with irb.for_range(
          0, out_shape, name="l"
      ) as l:  # pylint: disable=invalid-name
          for i in range(nrounds // 4):
              for j in range(nwords):
                  out_buf[out_offset + l * nwords + j] += key_schedule(
                      i, j
                  )  # wrapping
              for k in range(4):
                  for j in range(nwords // 2):
                      (
                          out_buf[out_offset + l * nwords + j * 2 + 0],
                          out_buf[out_offset + l * nwords + j * 2 + 1],
                      ) = mix(
                          out_buf[out_offset + l * nwords + j * 2 + 0],
                          out_buf[out_offset + l * nwords + j * 2 + 1],
                          _ROTATIONS[nwords][(i * 4 + k) % 8][j],
                      )
                  for j in range(nwords):
                      tmp[tmp_offset + l * nwords + j] = out_buf[
                          out_offset
                          + l * nwords
                          + _PERMUTATIONS[nwords][j]
                      ]
                  # number of rounds is even, so out always contains the result
                  (out_buf, tmp) = (tmp, out_buf)
                  (out_offset, tmp_offset) = (tmp_offset, out_offset)
  
file_path: tvm/python/tvm/topi/random/kernel.py
imports: import math
  import numpy
  import tvm
  import tvm.topi
  from ...None import tir
  from ...tir import ir_builder
label_desc: Use image node_id 1 for calling a global function with example usage: _threefry(irb, key_buf, key_offset, counter_buf, counter_offset, out_buf, out_offset, out_shape) and returns: list, full_key, unknown