# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # pylint: disable=invalid-name, too-many-nested-blocks """Backend kernels for cumsum operator.""" import math from tvm.script import tirx as T from tvm.tirx import PrimFunc def _is_power_of_two(n: int): """Check if n is a power of 2.""" return n > 0 and (n & (n - 1)) == 0 def gpu_2d_continuous_cumsum( ty_len: int = 4, tx_len: int = 32, thread_elem: int = 4, in_dtype: str = "int32", out_dtype: str | None = None, ) -> PrimFunc: """Generate GPU kernel for 2D continuous cumsum, i.e. The cumsum axis is -1 Parameters ---------- ty_len : int The length of `threadIdx.y` tx_len : int The length of `threadIdx.x` thread_elem : int The number of elements processed by single thread in_dtype : str The input data type out_dtype : Optional[str] The output data type, if None, it will be the same as in_dtype Returns ------- cumsum : PrimFunc The generated cumsum kernel """ out_dtype = out_dtype or in_dtype # Configuration for GPU kernel TX = T.int64(tx_len) # threadIdx.x TY = T.int64(ty_len) # threadIdx.y N = T.int64(thread_elem) # number of elements in single thread if not _is_power_of_two(TX) or not _is_power_of_two(TY) or not _is_power_of_two(N): raise ValueError("Configuration of TX, TY, N must be power of 2") # number of elements to be processed by single warp warp_elem = T.int64(tx_len * thread_elem) # number of elements to be processed by single block(SM) block_elem = T.int64(tx_len * ty_len * thread_elem) LOG_TX = T.int64(int(math.log2(tx_len))) LOG_BLOCK_N = T.int64(int(math.log2(tx_len * ty_len * thread_elem))) @T.macro def block_inclusive_inside_block( batch: T.int64, cur_len: T.int64, source: T.Buffer, output: T.Buffer, tmp_buf: T.Buffer, src_offset: T.int64, tmp_offset: T.int64, ): for by in T.thread_binding(batch, thread="blockIdx.y"): for bx in T.thread_binding(T.ceildiv(cur_len, block_elem), thread="blockIdx.x"): with T.sblock(): local_buf = T.sblock_alloc_buffer((thread_elem,), out_dtype, scope="local") shared_buf = T.sblock_alloc_buffer((block_elem,), out_dtype, scope="shared") for ty in T.thread_binding(TY, thread="threadIdx.y"): for tx in T.thread_binding(TX, thread="threadIdx.x"): tx_idx: T.let[T.int64] = ( bx * block_elem + ty * warp_elem + tx * thread_elem ) # Load data from global memory for i in T.vectorized(N): local_buf[i] = T.if_then_else( tx_idx + i < cur_len, T.Cast(out_dtype, source[by, src_offset + tx_idx + i]), T.Cast(out_dtype, 0), ) # Inclusive scan inside thread for i in T.unroll(1, N): local_buf[i] += local_buf[i - 1] # Store data to shared memory for i in T.vectorized(N): shared_buf[ty * warp_elem + tx * thread_elem + i] = local_buf[i] # Inclusive scan inside warp for i in T.unroll(LOG_TX): for j in T.vectorized(N): idx: T.let[T.int64] = ty * warp_elem + tx * thread_elem if tx >= (1 << i): shared_buf[idx + j] += shared_buf[ idx - (1 << i) * thread_elem + N - 1 ] # Inclusive scan inside block for i in T.unroll(1, TY): for j in T.vectorized(N): if ty == 0: idx: T.let[T.int64] = i * warp_elem + tx * thread_elem shared_buf[idx + j] += shared_buf[i * warp_elem - 1] # Write sum of block to global memory for i in T.vectorized(N): idx: T.let[T.int64] = ty * warp_elem + tx * thread_elem + i if bx * block_elem + idx < cur_len: output[by, src_offset + bx * block_elem + idx] = shared_buf[idx] if tx == 0 and ty == 0: for i in T.vectorized(N): tmp_buf[by, tmp_offset + bx] = shared_buf[block_elem - 1] @T.macro def update_cross_block( batch: T.int64, cur_len: T.int64, source: T.Buffer, output: T.Buffer, src_offset: T.int64, out_offset: T.int64, ): for by in T.thread_binding(batch, thread="blockIdx.y"): for bx in T.thread_binding(T.ceildiv(cur_len, block_elem), thread="blockIdx.x"): for ty in T.thread_binding(TY, thread="threadIdx.y"): for tx in T.thread_binding(TX, thread="threadIdx.x"): for i in T.serial(N): idx: T.let[T.int64] = bx * block_elem + ty * warp_elem + i * TX + tx if idx < cur_len: output[by, out_offset + idx] += T.if_then_else( bx > 0, source[by, src_offset + bx - 1], 0 ) @T.prim_func(private=True, s_tir=True) def cumsum(var_a: T.handle, var_out: T.handle): T.func_attr({"tirx.is_scheduled": True}) # prevent further scheduling m, n = T.int64(), T.int64() A = T.match_buffer(var_a, [m, n], dtype=in_dtype) Out = T.match_buffer(var_out, [m, n], dtype=out_dtype) Tmp = T.alloc_buffer([m, n], dtype=out_dtype) total_rounds: T.let[T.int64] = ( T.Cast("int64", T.ceil(T.log2(T.Cast("float32", n)))) // LOG_BLOCK_N ) block_inclusive_inside_block( m, n, A, Out, Tmp, src_offset=T.int64(0), tmp_offset=T.int64(0) ) for i in range(total_rounds): cur_len: T.let[T.int64] = T.ceildiv(n, 1 << (LOG_BLOCK_N * (i + 1))) block_inclusive_inside_block( m, cur_len, Tmp, Tmp, Tmp, src_offset=i * T.ceildiv(n, block_elem), tmp_offset=(i + 1) * T.ceildiv(n, block_elem), ) for i in range(total_rounds - 1): real_idx: T.let[T.int64] = total_rounds - 1 - i - 1 cur_len: T.let[T.int64] = T.ceildiv(n, 1 << (LOG_BLOCK_N * (real_idx + 1))) update_cross_block( m, cur_len, Tmp, Tmp, src_offset=(real_idx + 1) * T.ceildiv(n, block_elem), out_offset=real_idx * T.ceildiv(n, block_elem), ) update_cross_block(m, n, Tmp, Out, src_offset=0, out_offset=0) return cumsum