Add block construction function

cpp/include/raft/matrix/detail/matrix.cuh

@@ -197,6 +197,47 @@ void sliceMatrix(const m_t* in,
  slice<<<grid, block, 0, stream>>>(in, n_rows, n_cols, out, x1, y1, x2, y2);
 }
 
+/**
+ * @brief Kernel for copying a small matrix inside of a bigger matrix with a
+ * size matches that slice
+ * @param src_d: input matrix
+ * @param m: number of rows of input matrix
+ * @param n: number of columns of input matrix
+ * @param dst_d: output matrix
+ * @param x1, y1: coordinate of the top-left point of the wanted area (0-based)
+ * @param x2, y2: coordinate of the bottom-right point of the wanted area
+ * (1-based)
+ */
+template <typename m_t, typename idx_t = int>
+__global__ void slice_insert(
+ const m_t* src_d, idx_t n_rows, idx_t n_cols, m_t* dst_d, idx_t x1, idx_t y1, idx_t x2, idx_t y2)
+{
+ idx_t idx = threadIdx.x + blockDim.x * blockIdx.x;
+ idx_t dm = x2 - x1, dn = y2 - y1;
+ if (idx < dm * dn) {
+ idx_t i = idx % dm, j = idx / dm;
+ idx_t is = i + x1, js = j + y1;
+ dst_d[is + js * n_rows] = src_d[idx];
+ }
+}
+
+template <typename m_t, typename idx_t = int>
+void sliceMatrix_insert(const m_t* in,
+ idx_t n_rows,
+ idx_t n_cols,
+ m_t* out,
+ idx_t x1,
+ idx_t y1,
+ idx_t x2,
+ idx_t y2,
+ cudaStream_t stream)
+{
+ // Slicing
+ dim3 block(64);
+ dim3 grid(((x2 - x1) * (y2 - y1) + block.x - 1) / block.x);
+ slice_insert<<<grid, block, 0, stream>>>(in, n_rows, n_cols, out, x1, y1, x2, y2);
+}
+
 /**
  * @brief Kernel for copying the upper triangular part of a matrix to another
  * @param src: input matrix with a size of mxn

cpp/include/raft/matrix/slice.cuh

@@ -74,6 +74,40 @@ void slice(raft::device_resources const& handle,
  handle.get_stream());
 }
 
+/**
+ * @brief Insert a small matrix into a bigger matrix using a slice (in-place)
+ * @tparam m_t type of matrix elements
+ * @tparam idx_t integer type used for indexing
+ * @param[in] handle: raft handle
+ * @param[in] in: input matrix (column-major)
+ * @param[out] out: output matrix (column-major)
+ * @param[in] coords: coordinates of the insertion slice
+ * example: Slice the 2nd and 3rd columns of a 4x3 matrix: slice(handle, in, out, {0, 1, 4, 3});
+ */
+template <typename m_t, typename idx_t>
+void slice_insert(raft::device_resources const& handle,
+ raft::device_matrix_view<m_t, idx_t, col_major> in,
+ raft::device_matrix_view<m_t, idx_t, col_major> out,
+ slice_coordinates<idx_t> coords)
+{
+ RAFT_EXPECTS(coords.row2 > coords.row1, "row2 must be > row1");
+ RAFT_EXPECTS(coords.col2 > coords.col1, "col2 must be > col1");
+ RAFT_EXPECTS(coords.row1 >= 0, "row1 must be >= 0");
+ RAFT_EXPECTS(coords.row2 <= out.extent(0), "row2 must be <= number of rows in the output matrix");
+ RAFT_EXPECTS(coords.col1 >= 0, "col1 must be >= 0");
+ RAFT_EXPECTS(coords.col2 <= out.extent(1),
+ "col2 must be <= number of columns in the output matrix");
+
+ detail::sliceMatrix_insert(in.data_handle(),
+ out.extent(0),
+ out.extent(1),
+ out.data_handle(),
+ coords.row1,
+ coords.col1,
+ coords.row2,
+ coords.col2,
+ handle.get_stream());
+}
 /** @} */ // end group matrix_slice
 
 } // namespace raft::matrix