# I is a dense fixed iterator.
I = T.dense_fixed(m)
# J1 is a sparse fixed iterator, whose dependent iterator is I
# it has maximum length n and number of non-zero elements per row: c,
# the indices data are stored in the region started from indices_1 handle,
# and the index data type (in indices array) is int32.
J1 = T.sparse_fixed(I, (n, c), indices_1, idtype="int32")
# J2 is a dense variable iterator, whose dependent iterator is I,
# it has a maximum length of n,
# the indptr data are stored in the region started from indptr_2 handle,
# and the index data type (in indptr array) is int32.
J2 = T.dense_variable(I, n, indptr_2, idtype="int32")
# J3 is a sparse variable iterator, whose dependent iterator is J1,
# it has maximum length of n1, number of elements nnz in the space composed of (I, J1, J3),
# the indptr data are stored in the region started from indptr_3 handle,
# and the indices data are stored in the region started from indices_3 handle,
# the index data type (of indptr and indices array) is "int64")
J3 = T.sparse_variable(J1, (n1, nnz), (indptr_3, indices_3), idtype="int64")

I = T.dense_fixed(3)

J = T.sparse_variable(I, (4, 6), (indices_1=[1, 0, 2, 3, 1, 3], indptr=[0, 1, 4, 6]), idtype="int32")

A = T.match_sparse_buffer(a, (I, J), dtype="float32", scope="global")

with T.sp_iter([I, J, K], "SSR", "sddmm") as [i, j, k]:
    with T.init():
        Y[i, j] = 0.0
    Y[i, j] = Y[i, j] + A[i, k] * B[j, k] * X[i, j]

with T.sp_iter([I, J, K], "SSR", "sddmm") as [i, j, k]:
    with T.init():
        C[i, j] = 0.0
    C[i, j] = Y[i, j] + A[i, k] * B[j, k]

with T.init():
	C[vi, vbi, vf] = T.float32(0)

with T.init():
	C_data[vi * (blk * feat_size) + vbi * feat_size + vf] = T.float32(0)

"""
sparse
"""

class Axis: ...

def dense_fixed(length: PrimExpr, idtype: str = "int32"): ...
def dense_variable(
    parent_axis: Axis, 
    sizes: Tuple[Union[PrimExpr, int], Union[PrimExpr, int]],
    data: Var,
    idtype: str = "int32",
): ...
def sparse_fixed(
    parent_axis: Axis,
    sizes: Tuple[Union[PrimExpr, int], Union[PrimExpr, int]],
    data: Var,
    idtype: str = "int32",
): ...
def sparse_variable(
    parent_axis: Axis,
    sizes: Tuple[Union[PrimExpr, int], Union[PrimExpr, int]],
    data: Tuple[Var, Var],
    idtype: str = "int32",
): ...
def match_sparse_buffer(data: Var, axes: Sequence[Axis], dtype: str = "float32"): ...
def alloc_sparse_buffer(axes: Sequence[Axis], dtype: str, scope: str = "float32"): ...
def assume_buffer_domain(buf: Buffer, dom: Tuple[PrimExpr, PrimExpr]): ...

class iter(ContextManager):
    def __init__(
        self, axes: Sequence[Axis], iter_types: str, name_hint: str = ""
    ) -> None: ...
    def __enter__(self) -> Sequence[IterVar]: ...

@tvm._ffi.register_object("tir.sparse.Axis")
class Axis(Object):
    name: str
    parent: Optional["Axis"]
    length: PrimExpr
    nnz: PrimExpr
    nnz_cols: Optional[PrimExpr]
    indptr: Optional[Var]
    indices: Optional[Var]
    idtype: str
    sorted: bool

    def __init__(
        self, name, parent, length, nnz, nnz_cols, indptr, indices, idtype, sorted
    ) -> None:
        self.__init_handle_by_constructor__(_ffi_api.Axis, name, parent, length, nnz, nnz_cols, indptr, indices, idtype, sorted)  # type: ignore

class Node(Object):
    """Base class of all IR Nodes, implements astext function."""

    def astext(self, show_meta_data=True, annotate=None):
        return _ffi_api.AsText(self, show_meta_data, annotate)

    def __str__(self):
        return _ffi_api.PrettyPrint(self)

@tvm._ffi.register_object("tir.Var")
class Var(PrimExprWithOp):
    def __init__(self, name: str, dtype: Union[str, ir.Type], span: Optional[Span] = None):
        self.__init_handle_by_constructor__(_ffi_api.Var, name, dtype, span)  # type: ignore

@tvm._ffi.register_object("tir.sparse.FormatRewriteRule")
class FormatRewriteRule(Object):

    def __init__(
        self,
        name: str,
        new_format_desc: tvm.tir.PrimFunc,
        buffers_to_rewrite: List[str],
        axes_before_rewrite: List[str],
        axes_after_rewrite: List[str],
        axis_map: Dict[str, List[str]],
        idx_map: Union[Callable, IndexMap],
        inv_idx_map: Union[Callable, IndexMap],
    ) -> None:
        if isinstance(idx_map, Callable):
            idx_map = IndexMap.from_func(idx_map)
        if isinstance(inv_idx_map, Callable):
            inv_idx_map = IndexMap.from_func(inv_idx_map)
        self.__init_handle_by_constructor__(
            _ffi_api.FormatRewriteRule,
            name,
            new_format_desc,
            buffers_to_rewrite,
            axes_before_rewrite,
            axes_after_rewrite,
            axis_map,
            idx_map,
            inv_idx_map,
        )  # type: ignore

FormatRewriteRule(
    0 + "_" + 1,
    ell.specialize({nnz_cols_symbol: bucket_size}),
    ["A"],
    ["I", "J"],
    ["O", "I", "J"],
    {"I": ["O", "I"], "J": ["J"]},
    csr2ell_index_map,
    csr2ell_inv_index_map,
)

mod = tvm.IRModule.from_expr(csrmm)
mod = format_decompose(mod, rewrites)

Pass SparseFormatDecompose(Array<FormatRewriteRule> composable_formats,
                           bool include_format_rewrite_blks) {
  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
    return SparseFormatDecompose(std::move(composable_formats), std::move(f),
                                 include_format_rewrite_blks);
  };
  return CreatePrimFuncPass(pass_func, 0, "tir.SparseFormatDecompose", {});
}

def format_decompose(
    mod: IRModule,
    composable_formats: List["FormatRewriteRule"],
    include_format_rewrite_blks: bool = True,
):
    """Rewrite the sparse format of sparse buffers in the TIR scripts.

    Parameters
    ----------
    mod : IRModule
        The IRModule to lower.
    composable_formats : List[FormatRewriteRule]
        Composable formats is a list of rewrite rules.
    include_format_rewrite_blks : bool
        Whether to include format rewrite blocks in the output.
    """
    if not isinstance(mod, IRModule):
        raise TypeError("Expected IRModule, but got {}".format(type(mod)))
    return SparseFormatDecompose(composable_formats, include_format_rewrite_blks)(mod)

PrimFunc SparseFormxatDecompose(Array<FormatRewriteRule> composable_formats, PrimFunc f,
                                bool include_format_rewrite_blks = true) {
  CHECK(composable_formats.size() >= 1)
      << "The given composable formats length should be greater than or equal to 1.";
  // Only apply this pass to TIR that is not from TE schedules
  if (!IsFromLegacyTESchedule(f) && SparseTIRLevel(f) == 2) {
    // SparseFormatDecomposer rewriter(composable_formats);
    PrimFuncNode* fptr = f.CopyOnWrite();
    Array<PrimFunc> format_descs;
    Array<Axis> old_sp_axes = f->sp_axes;
    Array<SparseBuffer> old_buffers;
    for (const auto& kv : f->buffer_map) {
      old_buffers.push_back(Downcast<SparseBuffer>(kv.second));
    }
    for (const FormatRewriteRule& rule : composable_formats) {
      format_descs.push_back(AddSuffix(rule->new_format_desc, "_" + rule->name));
    }
    fptr->params = UpdateParams(format_descs, f->params);
    fptr->buffer_map = UpdateBufferMap(format_descs, f->buffer_map);
    fptr->sp_axes = UpdateSparseAxes(format_descs, f->sp_axes);
    Array<Stmt> format_rewrite_blks, compute_blks;
    // generate format rewrite blocks and compute blocks for each rule
    for (size_t i = 0; i < composable_formats.size(); ++i) {
      SparseFormatDecomposer rewriter(composable_formats[i], format_descs[i], old_sp_axes,
                                      old_buffers);
      rewriter(f->body);
      for (const Stmt& sp_iter : rewriter.format_rewrites_blks) {
        format_rewrite_blks.push_back(sp_iter);
      }
      for (const Stmt& sp_iter : rewriter.compute_blks) {
        compute_blks.push_back(sp_iter);
      }
    }
    // merge format rewrite and compute blocks.
    Array<Stmt> all_blks;
    if (include_format_rewrite_blks) {
      for (const Stmt& sp_iter : format_rewrite_blks) {
        all_blks.push_back(sp_iter);
      }
    }
    for (const Stmt& sp_iter : compute_blks) {
      all_blks.push_back(sp_iter);
    }
    Stmt new_body = all_blks.size() == 1 ? all_blks[0] : SeqStmt(all_blks);
    fptr->body = BlockRealize({}, const_true(), Block({}, {}, {}, "root", new_body, NullOpt, {}));
    // add composable flag.
    if (composable_formats.size() > 1) {
      Map<String, ObjectRef> new_attr_dict = fptr->attrs->dict;
      new_attr_dict.Set("composable", Integer(1));
      fptr->attrs = DictAttrs(new_attr_dict);
    }
    return f;
  } else {
    return f;
  }
}

def RemovePreprocess():
    # Remove the preprocess blocks/sparse iterations in the module.
    
    return _ffi_api.RemovePreprocess()  # type: ignore

PrimFunc RemovePreprocess(PrimFunc f) {
  if (!IsFromLegacyTESchedule(f)) {
    PrimFuncNode* fptr = f.CopyOnWrite();
    PreprocessRemover remover;
    fptr->body = remover(fptr->body);
    // insert extra parameters
    for (const Var& var : remover.extra_buffer_vars) {
      fptr->params.push_back(var);
      ICHECK(remover.extra_buffer_map.count(var))
          << "Internal error, extra_buffer_map do not have key " << var;
      fptr->buffer_map.Set(var, remover.extra_buffer_map.Get(var).value());
    }
    return f;
  } else {
    return f;
  }
}

Pass RemovePreprocess() {
  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
    return RemovePreprocess(std::move(f));
  };
  return CreatePrimFuncPass(pass_func, 0, "tir.RemovePreprocess", {});
}

    def sparse_reorder(self, block: SparseIterationRV, new_order: List[SpIterVar]) -> None:
        _ffi_api.ScheduleSparseReorder(  # type: ignore # pylint: disable=no-member
            self,
            block,
            new_order,
        )

    def sparse_fuse(self, block: SparseIterationRV, iters_to_fuse: List[SpIterVar]) -> None:
        _ffi_api.ScheduleSparseFuse(  # type: ignore # pylint: disable=no-member
            self,
            block,
            iters_to_fuse,
        )

sch = tir.Schedule(mod)
sp_iteration = sch.get_sparse_iteration("sddmm")
i, j, k = sch.get_sp_iters(sp_iteration)
sch.sparse_reorder(sp_iteration, [k, i, j])
sch.sparse_fuse(sp_iteration, [i, j])

SparseIteration SparseReorder(ScheduleState self, const SparseIteration& sp_iteration,
                              const Array<SpIterVar>& new_order) {
  // Step 1. Check whether the iterators in `new_order` are the same as `sp_iteration`'s iterators.
  CheckValidInputIterators(self, new_order, sp_iteration->sp_iter_vars);

  // Step 2. Check whether the new order does not break the iterator dependency.
  CheckDependency(self, new_order);

  // Step 3. Create the new SparseIteration.
  ObjectPtr<SparseIterationNode> p_new_sp_iteration =
      make_object<SparseIterationNode>(*sp_iteration.get());
  p_new_sp_iteration->sp_iter_vars = new_order;
  SparseIteration new_sp_iteration(p_new_sp_iteration);

  UpdateIRModule(self, sp_iteration, new_sp_iteration);
  return new_sp_iteration;
}

SparseIteration SparseFuse(ScheduleState self, const SparseIteration& sp_iteration,
                           const Array<SpIterVar>& iters_to_fuse) {
  // Step 1. Check match or not.
  int match_pos = CheckFuseMatch(self, iters_to_fuse, sp_iteration->sp_iter_vars);

  ObjectPtr<SparseIterationNode> p_new_sp_iteration =
      make_object<SparseIterationNode>(*sp_iteration.get());
  Array<SpIterVar> new_sp_iters;
  for (int i = 0; i < match_pos; ++i) {
    new_sp_iters.push_back(sp_iteration->sp_iter_vars[i]);
  }
  Array<Axis> axis_group;
  for (const SpIterVar& sp_iter_var : iters_to_fuse) {
    axis_group.push_back(sp_iter_var->axis);
  }
  for (size_t i = 0; i < iters_to_fuse.size(); ++i) {
    const SpIterVar& sp_iter_var = iters_to_fuse[i];
    Axis new_axis = FusedAxis(axis_group, i);
    new_sp_iters.push_back(SpIterVar(sp_iter_var->var, sp_iter_var->is_reduction, new_axis));
  }
  for (size_t i = match_pos + iters_to_fuse.size(); i < sp_iteration->sp_iter_vars.size(); ++i) {
    new_sp_iters.push_back(sp_iteration->sp_iter_vars[i]);
  }
  p_new_sp_iteration->sp_iter_vars = new_sp_iters;
  SparseIteration new_sp_iteration(p_new_sp_iteration);

  UpdateIRModule(self, sp_iteration, new_sp_iteration);
  return new_sp_iteration;
}

def lower_sparse_iter(mod: IRModule, check_invalid_binary_search: bool = False):
    """Lower sparse iterators in Sparse TIR.
    """
    
    if not isinstance(mod, IRModule):
        raise TypeError("Expected IRModule, but got {}".format(type(mod)))
    return LowerSparseIter(check_invalid_binary_search)(mod)

PrimFunc LowerSparseIter(PrimFunc f, bool check_invalid_binary_search) {
  // Only apply this pass to TIR that is not from TE schedules
  if (!IsFromLegacyTESchedule(f) && SparseTIRLevel(f) == 2) {
    PrimFuncNode* fptr = f.CopyOnWrite();
    // Step 1. Update the PrimFunc's buffer map.
    Map<Axis, SparseBuffer> axis_indptr_map, axis_indices_map;
    Array<BufferDomain> buf_doms;
    std::tie(axis_indptr_map, axis_indices_map, fptr->buffer_map, fptr->sp_axes, buf_doms) =
        UpdateMetadata(f);
    // Step 2. Lower iterations.
    IterTransformer lower_sparse(axis_indptr_map, axis_indices_map, fptr->sp_axes,
                                 check_invalid_binary_search);
    Stmt body = lower_sparse(std::move(fptr->body));
    // Step 3. Wrap with root block, insert bsearch blocks and allocated buffers.
    if (!lower_sparse.bsearch_structures.empty()) {
      Array<Stmt> seq;
      for (const auto& bsearch_struct : lower_sparse.bsearch_structures) {
        seq.push_back(bsearch_struct.body);
      }
      seq.push_back(body);
      body = SeqStmt(seq);
    }
    buf_doms = Concat(buf_doms, lower_sparse.alloc_buf_doms);
    Block root_block({}, {}, {}, "root", body, NullOpt, lower_sparse.root_alloc_buffers, {},
                     buf_doms);
    fptr->body = BlockRealize({}, const_true(), std::move(root_block));
    // Step 4. Lower sparse tir level.
    Map<String, ObjectRef> new_attr_dict = fptr->attrs->dict;
    new_attr_dict.Set("sparse_tir_level", Integer(1));
    fptr->attrs = DictAttrs(new_attr_dict);
    // Step 5. postprocess bufferload with possible invalid indices
    if (check_invalid_binary_search) {
      fptr->body = InvalidIndicesPostProcess()(std::move(fptr->body));
    }
    return f;
  } else {
    return f;
  }
}

Pass LowerSparseIter(bool check_invalid_binary_search) {
  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
    return LowerSparseIter(std::move(f), check_invalid_binary_search);
  };
  return CreatePrimFuncPass(pass_func, 0, "tir.LowerSparseIter", {});
}

sch = tir.Schedule(mod_preprocess)
blk = sch.get_block("binary_search_block_0_0")
(i,) = sch.get_loops(blk)
io, ii = sch.split(i, [None, 32])
sch.bind(ii, "threadIdx.x")
sch.bind(io, "blockIdx.x")

mod = lower_sparse_buffer(sch.mod)

PrimFunc LowerSparseBuffer(PrimFunc f) {
  // Only apply this pass to TIR that is not from TE schedules
  if (!IsFromLegacyTESchedule(f) && SparseTIRLevel(f) == 1) {
    bool is_horizontal_fuse = f->HasNonzeroAttr("horizontal_fuse");
    PrimFuncNode* fptr = f.CopyOnWrite();
    // Step 1. Update the PrimFunc's buffer map.
    fptr->buffer_map = std::move(UpdateBufferMap(f));
    // Step 2. Lower sparse buffers.
    fptr->body = BufferTransformer(fptr->sp_axes, fptr->buffer_map,
                                   is_horizontal_fuse)(std::move(fptr->body));
    // Step 3. Remove sparse axes
    fptr->sp_axes.clear();
    // Step 4. Lower sparse tir level
    Map<String, ObjectRef> new_attr_dict = fptr->attrs->dict;
    new_attr_dict.Set("sparse_tir_level", Integer(0));
    fptr->attrs = DictAttrs(new_attr_dict);
    return f;
  } else {
    return f;
  }
}

namespace transform {

/*!
 * \brief The lowering pass from TIR to Sparse TIR.
 */
Pass LowerSparseBuffer() {
  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
    return LowerSparseBuffer(std::move(f));
  };
  return CreatePrimFuncPass(pass_func, 0, "tir.LowerSparseBuffer", {});
}

f = tvm.build(mod, target="cuda")

evaluator = f.time_evaluator(f.entry_name, tvm.cuda(0), number=100)

import dgl
import tvm
import tvm.testing
import tvm.tir as tir
import scipy.sparse as sp
import argparse
import numpy as np
import torch as th
from tvm.script import tir as T
from tvm.sparse import (
    FormatRewriteRule,
    lower_sparse_buffer,
    lower_sparse_iter,
    column_part_hyb,
    format_decompose,
)
import tvm.sparse
from utils import get_dataset, ell


@T.prim_func
def csrmm(
    a: T.handle,
    b: T.handle,
    c: T.handle,
    indptr: T.handle,
    indices: T.handle,
    m: T.int32,
    n: T.int32,
    num_tiles: T.int32,
    nnz: T.int32,
    cwm: T.int32,
) -> None:
    T.func_attr({"global_symbol": "main", "tir.noalias": True, "sparse_tir_level": 2})
    I = T.dense_fixed(m)
    J = T.sparse_variable(I, (n, nnz), (indptr, indices), "int32")
    J_detach = T.dense_fixed(n)
    K1 = T.dense_fixed(num_tiles)
    K2 = T.dense_fixed(cwm)
    K3 = T.dense_fixed(32)
    A = T.match_sparse_buffer(a, (I, J), "float32")
    B = T.match_sparse_buffer(b, (J_detach, K1, K2, K3), "float32")
    C = T.match_sparse_buffer(c, (I, K1, K2, K3), "float32")
    with T.sp_iter([I, J, K1, K2, K3], "SRSSS", "csrmm") as [i, j, k1, k2, k3]:
        with T.init():
            C[i, k1, k2, k3] = 0.0
        C[i, k1, k2, k3] = C[i, k1, k2, k3] + A[i, j] * B[j, k1, k2, k3]


class TorchOpTimer(object):
    def __enter__(self):
        self.start_event = th.cuda.Event(enable_timing=True)
        self.end_event = th.cuda.Event(enable_timing=True)
        self.start_event.record()
        return self

    def __exit__(self, type, value, traceback):
        self.end_event.record()
        th.cuda.synchronize()  # Wait for the events to be recorded!
        self.time = self.start_event.elapsed_time(self.end_event) / 1e3


def csr2ell_inv_index_map(o, i, j):
    return i, j


def csr2ell_index_map(i, j):
    return 0, i, j


cached_bucketing_format = None


def bench_hyb(
    g,
    x,
    y_golden,
    feat_size=128,
    bucket_sizes=[],
    coersening_factor=2,
    num_col_parts=1,
    use_implicit_unroll=False,
):
    num_buckets = len(bucket_sizes)
    coersening_factor = min(coersening_factor, feat_size // 32)
    indptr, indices, _ = g.adj_sparse("csc")
    m = g.num_dst_nodes()
    n = g.num_src_nodes()
    nnz = g.num_edges()
    global cached_bucketing_format
    if cached_bucketing_format is None:
        indptr_nd = tvm.nd.array(indptr.numpy(), device=tvm.cpu())
        indices_nd = tvm.nd.array(indices.numpy(), device=tvm.cpu())
        cached_bucketing_format = column_part_hyb(
            m, n, indptr_nd, indices_nd, num_col_parts, bucket_sizes
        )
    row_indices, col_indices, mask = cached_bucketing_format

    # rewrite csrmm
    nnz_cols_symbol = ell.params[-1]
    rewrites = []
    for part_id in range(num_col_parts):
        for bucket_id, bucket_size in enumerate(bucket_sizes):
            rewrites.append(
                FormatRewriteRule(
                    str(part_id) + "_" + str(bucket_id),
                    ell.specialize({nnz_cols_symbol: bucket_size}),
                    ["A"],
                    ["I", "J"],
                    ["O", "I", "J"],
                    {"I": ["O", "I"], "J": ["J"]},
                    csr2ell_index_map,
                    csr2ell_inv_index_map,
                )
            )
    mod = tvm.IRModule.from_expr(csrmm)
    mod = format_decompose(mod, rewrites)
    mod = tvm.tir.transform.RemovePreprocess()(mod)

    # specialize
    params = mod["main"].params
    param_map = {
        params[5]: m,  # m
        params[6]: n,  # n
        params[7]: feat_size // coersening_factor // 32,  # num_tiles,
        params[8]: nnz,  # nnz
        params[9]: coersening_factor,  # coersening_factor
    }
    for part_id in range(num_col_parts):
        for bucket_id in range(num_buckets):
            param_map[params[10 + 7 * (part_id * num_buckets + bucket_id) + 4]] = m
            param_map[params[10 + 7 * (part_id * num_buckets + bucket_id) + 5]] = n
            param_map[params[10 + 7 * (part_id * num_buckets + bucket_id) + 6]] = row_indices[
                part_id
            ][bucket_id].shape[0]

    mod["main"] = mod["main"].specialize(param_map).with_attr("horizontal_fuse", True)

    # schedule
    sch = tvm.tir.Schedule(mod)
    for sp_iter_name in [
        "csrmm_{}_{}".format(i, j) for j in range(num_buckets) for i in range(num_col_parts)
    ]:
        sp_iteration = sch.get_sparse_iteration(sp_iter_name)
        o, i, j, k1, k2, k3 = sch.get_sp_iters(sp_iteration)
        sch.sparse_fuse(sp_iteration, [o, i])

    mod = sch.mod
    mod = tvm.sparse.lower_sparse_iter(mod)
    sch = tvm.tir.Schedule(mod)
    for part_id in range(num_col_parts):
        for bucket_id, bucket_size in enumerate(bucket_sizes):
            is_atomic = num_col_parts > 1 or bucket_id + 1 == num_buckets
            blk = sch.get_block("csrmm_{}_{}0".format(part_id, bucket_id))
            i, j, foo, foi, fi = sch.get_loops(blk)
            sch.reorder(foo, fi, j, foi)
            if is_atomic:
                sch.annotate(blk, "atomic", True)
                write_blk = sch.reverse_cache_write(blk, 0, "local")
                sch.reverse_compute_at(write_blk, fi, True)
                # sch.unroll(sch.get_loops(write_blk)[-2])
            sch.bind(fi, "threadIdx.x")
            sch.bind(foo, "blockIdx.y")
            sch.unroll(foi)
            if use_implicit_unroll:
                sch.annotate(foi, "pragma_unroll_explicit", 0)
            sch.unroll(j)
            if use_implicit_unroll:
                sch.annotate(j, "pragma_unroll_explicit", 0)
            io, ioi, ii = sch.split(i, [None, bucket_sizes[-1] // bucket_size, 8])
            sch.bind(io, "blockIdx.x")
            sch.bind(ii, "threadIdx.y")
            init_blk = sch.decompose_reduction(blk, fi)
            ax0, ax1 = sch.get_loops(init_blk)[-2:]
            sch.bind(ax0, "threadIdx.x")
            sch.unroll(ax1)
            if use_implicit_unroll:
                sch.annotate(ax1, "pragma_unroll_explicit", 0)

    mod = tvm.sparse.lower_sparse_buffer(sch.mod)
    mod = tvm.tir.transform.RemoveUnusedArgs()(mod)
    f = tvm.build(mod, target="cuda")

    # prepare nd array
    b_nd = tvm.nd.array(
        x.numpy().reshape(-1).astype("float32"),
        device=tvm.cuda(0),
    )
    c_nd = tvm.nd.array(np.zeros((n * feat_size,)).astype("float32"), device=tvm.cuda(0))
    # prepare args
    args = [b_nd, c_nd]

    for part_id in range(num_col_parts):
        for bucket_id, _ in enumerate(bucket_sizes):
            weight = tvm.nd.array(
                mask[part_id][bucket_id].numpy().reshape(-1).astype("float32"), device=tvm.cuda(0)
            )
            rows = tvm.nd.array(
                row_indices[part_id][bucket_id].numpy().astype("int32"), device=tvm.cuda(0)
            )
            cols = tvm.nd.array(
                col_indices[part_id][bucket_id].numpy().reshape(-1).astype("int32"),
                device=tvm.cuda(0),
            )
            args += [weight, rows, cols]

    # test accuracy
    f(*args)
    tvm.testing.assert_allclose(c_nd.numpy().reshape(-1, feat_size), y_golden.numpy(), rtol=1e-4)

    # evaluate time
    evaluator = f.time_evaluator(f.entry_name, tvm.cuda(0), number=100)
    print("tir hyb time: {:.5f}ms".format(evaluator(*args).mean * 1000))


col_part_config = {
    "arxiv": 1,
    "proteins": 8,
    "pubmed": 1,
    "citeseer": 1,
    "cora": 1,
    "ppi": 16,
    "reddit": 8,
    "products": 16,
}

bucketing_config = {
    "arxiv": [1, 2, 4, 8, 16, 32],
    "proteins": [1, 2, 4, 8, 16, 32, 64, 128, 256],
    "pubmed": [1, 2, 4, 8, 16, 32],
    "citeseer": [1, 2, 4],
    "cora": [1, 2, 4],
    "ppi": [1, 2, 4, 8, 16, 32],
    "products": [1, 2, 4, 8, 16, 32],
    "reddit": [1, 2, 4, 8, 16, 32, 64, 128, 256, 512],
}


if __name__ == "__main__":
    parser = argparse.ArgumentParser("hybrid format spmm in sparse-tir")
    parser.add_argument("--dataset", "-d", type=str, default="arxiv", help="dataset name")
    parser.add_argument("--implicit-unroll", "-i", action="store_true", help="use implicit unroll")
    args = parser.parse_args()
    name = args.dataset
    g = get_dataset(name)

    for feat_size in [32, 64, 128, 256, 512]:
        print("feat_size =", feat_size)
        x = th.rand((g.num_src_nodes(), feat_size))
        y_golden = dgl.ops.copy_u_sum(g, x)
        bench_hyb(
            g,
            x,
            y_golden,
            feat_size=feat_size,
            bucket_sizes=bucketing_config[name],
            coersening_factor=2,
            num_col_parts=col_part_config[name],
            use_implicit_unroll=args.implicit_unroll,
        )