""" The compute function and schedules for SDDMM kernels written in TVM. """
import tvm
from tvm import te


def sddmm_tree_reduction_gpu(idx_type, feat_type):
    """ SDDMM kernels on GPU optimized with Tree Reduction.
    
    Parameters
    ----------
    idx_type : str
        The data type for indexing tensors.
    feat_type : str
        The data type of feature tensor.

    Returns
    -------
    IRModule
        The result IRModule.
    """
    # define vars and placeholders
    nnz = te.var('nnz', idx_type)
    num_rows = te.var('num_rows', idx_type)
    num_cols = te.var('num_cols', idx_type)
    H = te.var('num_heads', idx_type)
    D = te.var('feat_len', idx_type)
    row = te.placeholder((nnz,), idx_type, 'row')
    col = te.placeholder((nnz,), idx_type, 'col')
    ufeat = te.placeholder((num_rows, H, D), feat_type, 'ufeat')
    vfeat = te.placeholder((num_cols, H, D), feat_type, 'vfeat')
    # define edge computation function
    def edge_func(eid, h, i):
        k = te.reduce_axis((0, D), name='k')
        return te.sum(ufeat[row[eid], h, k] * vfeat[col[eid], h, k], axis=k)
    out = te.compute((nnz, H, tvm.tir.IntImm(idx_type, 1)), edge_func, name='out')
    # define schedules
    sched = te.create_schedule(out.op)
    edge_axis, head_axis, _ = out.op.axis
    reduce_axis = out.op.reduce_axis[0]
    _, red_inner = sched[out].split(reduce_axis, factor=32)
    edge_outer, edge_inner = sched[out].split(edge_axis, factor=32)
    sched[out].bind(red_inner, te.thread_axis('threadIdx.x'))
    sched[out].bind(edge_inner, te.thread_axis('threadIdx.y'))
    sched[out].bind(edge_outer, te.thread_axis('blockIdx.x'))
    sched[out].bind(head_axis, te.thread_axis('blockIdx.y'))
    return tvm.lower(sched, [row, col, ufeat, vfeat, out],
                     name='SDDMMTreeReduction_{}_{}'.format(idx_type, feat_type))


if __name__ == '__main__':
    kernel0 = sddmm_tree_reduction_gpu('int32', 'float32')
    print(kernel0)