[AdvancedCompiler]Cat(cpp wrapper)

ctests/CMakeLists.txt

@@ -22,6 +22,10 @@ target_link_libraries(test_triton_rope
   PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)
 add_test(NAME test_triton_rope COMMAND test_triton_rope)
 
+add_executable(test_triton_cat test_triton_cat.cpp)
+target_link_libraries(test_triton_cat
+  PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)
+add_test(NAME test_triton_cat COMMAND test_triton_cat)
 add_executable(test_triton_bmm test_triton_bmm.cpp)
 target_link_libraries(test_triton_bmm
   PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)

ctests/test_triton_cat.cpp

@@ -0,0 +1,133 @@
+#include "flag_gems/operators.h"
+#include "gtest/gtest.h"
+#include "torch/torch.h"
+
+TEST(TritonCatTest, basictest) {
+  const torch::Device device(torch::kCUDA, 0);
+  torch::Tensor t1 = torch::randn({2, 3}, device);
+  torch::Tensor t2 = torch::randn({4, 3}, device);
+
+  torch::Tensor out_torch = torch::cat({t1, t2}, 0);
+  torch::Tensor out_gems = flag_gems::cat({t1, t2}, 0);
+
+  EXPECT_TRUE(torch::equal(out_torch, out_gems));
+}
+
+TEST(TritonCatTest, 2dimtest) {
+  const torch::Device device(torch::kCUDA, 0);
+  torch::Tensor t1 = torch::randn({3, 2}, device);
+  torch::Tensor t2 = torch::randn({3, 4}, device);
+
+  int dim_to_test = 1;
+  torch::Tensor out_torch = torch::cat({t1, t2}, dim_to_test);
+  torch::Tensor out_gems = flag_gems::cat({t1, t2}, dim_to_test);
+
+  EXPECT_TRUE(torch::equal(out_gems, out_torch));
+
+  EXPECT_EQ(out_gems.size(0), 3);
+  EXPECT_EQ(out_gems.size(1), 6);
+}
+
+TEST(TritonCatTest, 3dimtest) {
+  const torch::Device device(torch::kCUDA, 0);
+
+  torch::Tensor t1 = torch::randn({3, 2, 4}, device);
+  torch::Tensor t2 = torch::randn({3, 5, 4}, device);
+
+  int dim_to_test = 1;
+
+  torch::Tensor out_torch = torch::cat({t1, t2}, dim_to_test);
+  torch::Tensor out_gems = flag_gems::cat({t1, t2}, dim_to_test);
+
+  EXPECT_TRUE(torch::equal(out_torch, out_gems));
+}
+
+TEST(TritonCatTest, 4dimtest) {
+  const torch::Device device(torch::kCUDA, 0);
+  auto options = torch::TensorOptions().device(device).dtype(torch::kFloat32);
+
+  torch::Tensor t1 = torch::randn({2, 3, 4, 5}, options);
+  torch::Tensor t2 = torch::randn({2, 6, 4, 5}, options);
+
+  int dim_to_test = 1;
+  torch::Tensor out_torch = torch::cat({t1, t2}, dim_to_test);
+  torch::Tensor out_gems = flag_gems::cat({t1, t2}, dim_to_test);
+
+  EXPECT_TRUE(torch::equal(out_torch, out_gems));
+}
+
+TEST(TritonCatTest, IntegerConcatenation) {
+  const torch::Device device(torch::kCUDA, 0);
+  auto options = torch::TensorOptions().device(device).dtype(torch::kInt32);
+
+  torch::Tensor t1 = torch::randint(0, 100, {2, 3, 4}, options);
+  torch::Tensor t2 = torch::randint(0, 100, {2, 3, 4}, options);
+
+  int dim_to_test = 2;
+  torch::Tensor out_torch = torch::cat({t1, t2}, dim_to_test);
+  torch::Tensor out_gems = flag_gems::cat({t1, t2}, dim_to_test);
+
+  EXPECT_TRUE(torch::equal(out_torch, out_gems));
+}
+
+TEST(TritonCatTest, EmptyTensorConcatenation) {
+  const torch::Device device(torch::kCUDA, 0);
+  auto options = torch::TensorOptions().device(device).dtype(torch::kFloat32);
+
+  torch::Tensor t1 = torch::randn({0, 3}, options);
+  torch::Tensor t2 = torch::randn({2, 3}, options);
+
+  torch::Tensor out_torch = torch::cat({t1, t2}, 0);
+  torch::Tensor out_gems = flag_gems::cat({t1, t2}, 0);
+
+  EXPECT_TRUE(torch::equal(out_torch, out_gems));
+
+  torch::Tensor t3 = torch::randn({0, 3}, options);
+  torch::Tensor t4 = torch::randn({0, 3}, options);
+
+  torch::Tensor out_torch_both_empty = torch::cat({t3, t4}, 0);
+  torch::Tensor out_gems_both_empty = flag_gems::cat({t3, t4}, 0);
+
+  EXPECT_TRUE(torch::equal(out_torch_both_empty, out_gems_both_empty));
+  EXPECT_EQ(out_gems_both_empty.numel(), 0);
+}
+
+TEST(TritonCatTest, 3tensorcat) {
+  const torch::Device device(torch::kCUDA, 0);
+  auto options = torch::TensorOptions().device(device).dtype(torch::kFloat32);
+
+  torch::Tensor t1 = torch::randn({2, 3}, options);
+  torch::Tensor t2 = torch::randn({4, 3}, options);
+  torch::Tensor t3 = torch::randn({1, 3}, options);
+
+  int dim_to_test = 0;
+  torch::Tensor out_torch = torch::cat({t1, t2, t3}, dim_to_test);
+  torch::Tensor out_gems = flag_gems::cat({t1, t2, t3}, dim_to_test);
+
+  EXPECT_TRUE(torch::equal(out_torch, out_gems));
+  EXPECT_EQ(out_gems.size(0), 7);
+  EXPECT_EQ(out_gems.size(1), 3);
+}
+
+TEST(TritonCatTest, HandlesNonContiguousInput) {
+  const torch::Device device(torch::kCUDA, 0);
+  auto options = torch::TensorOptions().device(device).dtype(torch::kFloat32);
+
+  torch::Tensor t_base = torch::randn({2, 3, 4}, options);
+
+  torch::Tensor t1_non_contiguous = t_base.transpose(1, 2);
+
+  torch::Tensor t2_contiguous = torch::randn({2, 4, 5}, options);
+
+  ASSERT_FALSE(t1_non_contiguous.is_contiguous());
+
+  int dim_to_test = 2;
+  torch::Tensor out_torch = torch::cat({t1_non_contiguous, t2_contiguous}, dim_to_test);
+  torch::Tensor out_gems = flag_gems::cat({t1_non_contiguous, t2_contiguous}, dim_to_test);
+
+  EXPECT_TRUE(torch::equal(out_torch, out_gems));
+
+  EXPECT_EQ(out_gems.size(0), 2);
+  EXPECT_EQ(out_gems.size(1), 4);
+  EXPECT_EQ(out_gems.size(2), 8);
+}

include/flag_gems/operators.h

@@ -30,5 +30,6 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding(
     const std::optional<at::Tensor> &position_ids = std::nullopt,
     bool rotary_interleaved = false);
 
+at::Tensor cat(const at::TensorList &tensors, int64_t dim = 0);
 at::Tensor bmm(const at::Tensor &A, const at::Tensor &B);
 }  // namespace flag_gems

lib/CMakeLists.txt

@@ -7,6 +7,7 @@ add_library(operators
             rms_norm.cpp
             fused_add_rms_norm.cpp
             rotary_embedding.cpp
+            cat.cpp
             bmm.cpp)
 target_include_directories(operators
   PUBLIC

lib/cat.cpp

@@ -0,0 +1,110 @@
+#include "c10/cuda/CUDAStream.h"
+#include "flag_gems/operators.h"
+#include "flag_gems/utils.h"
+#include "triton_jit/triton_jit_function.h"
+
+namespace flag_gems {
+using namespace triton_jit;
+
+at::Tensor cat(const at::TensorList& tensors, int64_t dim) {
+  TORCH_CHECK(tensors.size() > 0, "torch.cat(): expected a non-empty list of Tensors");
+  if (tensors.size() == 1) {
+    return tensors[0];
+  }
+  const auto& first_tensor = tensors[0];
+  int64_t ndim = first_tensor.dim();
+  TORCH_CHECK(dim >= -ndim && dim < ndim, "cat(): dimension out of range");
+  if (dim < 0) {
+    dim += ndim;
+  }
+  const at::IntArrayRef first_shape = first_tensor.sizes();
+  for (size_t i = 1; i < tensors.size(); ++i) {
+    const auto& current_tensor = tensors[i];
+    TORCH_CHECK(current_tensor.dim() == ndim,
+                "Tensors must have same number of dimensions: got ",
+                ndim,
+                " and ",
+                current_tensor.dim());
+    const at::IntArrayRef current_shape = current_tensor.sizes();
+    for (int64_t d = 0; d < ndim; ++d) {
+      if (d == dim) continue;
+      TORCH_CHECK(current_shape[d] == first_shape[d],
+                  "Sizes of tensors must match except in dimension ",
+                  dim,
+                  ". Expected size ",
+                  first_shape[d],
+                  " but got size ",
+                  current_shape[d],
+                  " for tensor number ",
+                  i);
+    }
+  }
+
+  std::vector<int64_t> out_shape_vec = first_shape.vec();
+  int64_t cat_dim_size = 0;
+  for (const auto& t : tensors) {
+    cat_dim_size += t.size(dim);
+  }
+  out_shape_vec[dim] = cat_dim_size;
+  at::Tensor out = at::empty(out_shape_vec, first_tensor.options());
+
+  std::vector<int64_t> storage_offsets;
+  int64_t current_storage_offset = 0;
+  storage_offsets.push_back(current_storage_offset);
+  int64_t out_stride_for_dim = out.stride(dim);
+  for (size_t i = 0; i < tensors.size() - 1; ++i) {
+    current_storage_offset += tensors[i].size(dim) * out_stride_for_dim;
+    storage_offsets.push_back(current_storage_offset);
+  }
+
+  const TritonJITFunction copy_kernel_func =
+      TritonJITFunction::getInstance(std::string(utils::get_triton_src_path() / "cat_copy.py"),
+                                     "strided_copy_kernel");
+  c10::DeviceGuard guard(out.device());
+  c10::cuda::CUDAStream stream = c10::cuda::getCurrentCUDAStream();
+  CUstream raw_stream = static_cast<CUstream>(stream.stream());
+
+  for (size_t i = 0; i < tensors.size(); ++i) {
+    const auto& input_tensor = tensors[i];
+    if (input_tensor.numel() == 0) continue;
+
+    at::Tensor output_view = at::as_strided(out, input_tensor.sizes(), out.strides(), storage_offsets[i]);
+
+    auto options = torch::TensorOptions().device(input_tensor.device()).dtype(torch::kInt64);
+    at::Tensor in_strides = torch::tensor(input_tensor.strides(), options);
+    at::Tensor out_strides = torch::tensor(output_view.strides(), options);
+    at::Tensor shapes = torch::tensor(input_tensor.sizes(), options);
+
+    int64_t ndim_val = input_tensor.dim();
+    int64_t num_elements = input_tensor.numel();
+
+    constexpr int BLOCK_SIZE = 256;
+    constexpr int MAX_DIMS = 8;
+    TORCH_CHECK(ndim_val <= MAX_DIMS,
+                "Tensor dimension ",
+                ndim_val,
+                " exceeds the maximum supported by the kernel (",
+                MAX_DIMS,
+                ")");
+
+    unsigned int grid_x = (num_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;
+
+    copy_kernel_func(raw_stream,
+                     grid_x,
+                     1,
+                     1,
+                     4,
+                     2,
+                     input_tensor,
+                     output_view,
+                     in_strides,
+                     out_strides,
+                     shapes,
+                     ndim_val,
+                     num_elements,
+                     BLOCK_SIZE,
+                     MAX_DIMS);
+  }
+  return out;
+}
+}  // namespace flag_gems

src/flag_gems/csrc/cstub.cpp

@@ -20,6 +20,7 @@ TORCH_LIBRARY(flag_gems, m) {
   m.def(
       "rotary_embedding(Tensor q, Tensor k, Tensor cos, Tensor sin, Tensor? position_ids=None, "
       "bool rotary_interleaved=False) -> (Tensor, Tensor)");  // q and k may be view to other size
+  m.def("cat(Tensor[] tensors, int dim=0) -> Tensor");
   m.def("bmm(Tensor self, Tensor mat2) -> Tensor");
 }
 
@@ -32,6 +33,7 @@ TORCH_LIBRARY_IMPL(flag_gems, CUDA, m) {
   // Rotary embedding
   m.impl("rotary_embedding", TORCH_FN(rotary_embedding));
   m.impl("rotary_embedding_inplace", TORCH_FN(rotary_embedding_inplace));
+  m.impl("cat", TORCH_FN(cat));
   m.impl("bmm", TORCH_FN(bmm));
 }
 }  // namespace flag_gems

triton_src/cat_copy.py

@@ -0,0 +1,40 @@
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def strided_copy_kernel(
+    in_ptr,
+    out_ptr,
+    in_strides_ptr,
+    out_strides_ptr,
+    shapes_ptr,
+    ndim,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+    MAX_DIMS: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+
+    offsets = offsets.to(tl.int64)
+    mask = offsets < n_elements
+
+    remaining_offset = offsets
+    in_physical_offset = tl.zeros_like(offsets)
+    out_physical_offset = tl.zeros_like(offsets)
+
+    for i in range(MAX_DIMS - 1, -1, -1):
+        is_real_dim = i < ndim
+        current_shape = tl.load(shapes_ptr + i, mask=is_real_dim, other=1)
+        in_stride_val = tl.load(in_strides_ptr + i, mask=is_real_dim, other=0)
+        out_stride_val = tl.load(out_strides_ptr + i, mask=is_real_dim, other=0)
+
+        current_index = remaining_offset % current_shape
+        remaining_offset = remaining_offset // current_shape
+
+        in_physical_offset += current_index * in_stride_val
+        out_physical_offset += current_index * out_stride_val
+
+    x = tl.load(in_ptr + in_physical_offset, mask=mask)
+    tl.store(out_ptr + out_physical_offset, x, mask=mask)