[FEA] Support for more general use of COO with Index_Type other than int

[jinsolp]

Description

Currently the COO class has Index_Type=int as default. Declared like below in raft/cpp/include/raft/sparse/detail/coo.cuh

template <typename T, typename Index_Type = int>
class COO {
 protected:
  rmm::device_uvector<Index_Type> rows_arr;
  rmm::device_uvector<Index_Type> cols_arr;
  rmm::device_uvector<T> vals_arr;

 public:
  Index_Type nnz;
  Index_Type n_rows;
  Index_Type n_cols;

However, most of the functions in raft that have COO as its argument is declared in a way that doesn't really allow other Index_Type types.

For example, coo_symmetrize in raft/cpp/include/raft/sparse/linalg/detail/symmetrize.cuh is declared like this

template <int TPB_X = 128, typename T, typename Lambda>
void coo_symmetrize(COO<T>* in,
                    COO<T>* out,
                    Lambda reduction_op,
                    cudaStream_t stream)

So we can actually only use COO that has Index_Type int.

TODO

1. change the setSize() and allocate() functions in COO to take Index_Types instead of int

// original
  void setSize(int n_rows, int n_cols)
  {
    this->n_rows = n_rows;
    this->n_cols = n_cols;
  }

// should change to...
 void setSize(Index_Type n_rows, Index_Type n_cols)
  {
    this->n_rows = n_rows;
    this->n_cols = n_cols;
  }

2. change function signatures and add template arguments for functions that have COO as its argument.

// original
template <int TPB_X = 128, typename T, typename Lambda>
void coo_symmetrize(COO<T>* in,
                    COO<T>* out,
                    Lambda reduction_op,  // two-argument reducer
                    cudaStream_t stream)
// should change to...
template <int TPB_X = 128, typename T, typename Lambda, typename Index_Type=int>
void coo_symmetrize(COO<T, Index_Type>* in,
                    COO<T, Index_Type>* out,
                    Lambda reduction_op,  // two-argument reducer
                    cudaStream_t stream)