[AdvancedCompiler]bmm(cpp wrapper)

ctests/CMakeLists.txt

@@ -21,3 +21,8 @@ add_executable(test_triton_rope test_triton_rope.cpp)
 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_bmm test_triton_bmm.cpp)
+target_link_libraries(test_triton_bmm
+  PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)
+add_test(NAME test_triton_bmm COMMAND test_triton_bmm)

ctests/test_triton_bmm.cpp

@@ -0,0 +1,17 @@
+#include <gtest/gtest.h>
+#include "flag_gems/operators.h"
+#include "torch/torch.h"
+
+TEST(blas_op_test, bmm) {
+  const torch::Device device(torch::kCUDA, 0);
+
+  const int B = 5, M = 10, K = 10, N = 10;
+
+  torch::Tensor batch1 = torch::randn({B, M, K}, device);
+  torch::Tensor batch2 = torch::randn({B, K, N}, device);
+
+  torch::Tensor out_torch = flag_gems::bmm(batch1, batch2);
+  torch::Tensor out_triton = flag_gems::bmm(batch1, batch2);
+
+  EXPECT_TRUE(torch::allclose(out_torch, out_triton));
+}

include/flag_gems/operators.h

@@ -29,4 +29,6 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding(
     const at::Tensor &sin,
     const std::optional<at::Tensor> &position_ids = std::nullopt,
     bool rotary_interleaved = false);
+
+at::Tensor bmm(const at::Tensor &A, const at::Tensor &B);
 }  // namespace flag_gems

lib/CMakeLists.txt

@@ -6,7 +6,8 @@ add_library(operators
             mm.cpp
             rms_norm.cpp
             fused_add_rms_norm.cpp
-            rotary_embedding.cpp)
+            rotary_embedding.cpp
+            bmm.cpp)
 target_include_directories(operators
   PUBLIC
     $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}/include>

lib/bmm.cpp

@@ -0,0 +1,70 @@
+#include "flag_gems/operators.h"
+#include "flag_gems/utils.h"
+
+#include <iostream>
+#include "c10/cuda/CUDAStream.h"
+#include "triton_jit/triton_jit_function.h"
+
+namespace flag_gems {
+using namespace triton_jit;
+
+at::Tensor bmm(const at::Tensor& A, const at::Tensor& B) {
+  TORCH_CHECK(A.dim() == 3 && B.dim() == 3, "both the tensors must be 3-D");
+  TORCH_CHECK(A.dtype() == B.dtype(),
+              "expected a and b to have the same dtype, but got: ",
+              A.dtype(),
+              " != ",
+              B.dtype());
+  at::IntArrayRef A_sizes = A.sizes();
+  at::IntArrayRef B_sizes = B.sizes();
+
+  at::Tensor A_contig = A.contiguous();
+  at::Tensor B_contig = B.contiguous();
+
+  at::Tensor out = at::empty({A_sizes[0], A_sizes[1], B_sizes[2]}, A_contig.options());
+
+  const TritonJITFunction& f =
+      TritonJITFunction::getInstance(std::string(utils::get_triton_src_path() / "bmm.py"), "bmm_kernel");
+
+  c10::DeviceGuard guard(out.device());
+  c10::cuda::CUDAStream stream = c10::cuda::getCurrentCUDAStream();
+  CUstream raw_stream = static_cast<CUstream>(stream.stream());
+  const int GROUP_M = 8;
+  const int TILE_M = 128;
+  const int TILE_N = 128;
+  const int TILE_K = 32;
+  const int M = A_sizes[1];
+  const int N = B_sizes[2];
+  const int K = A_sizes[2];
+  unsigned int grid_x = (M + (TILE_M - 1)) / TILE_M;
+  unsigned int grid_y = (N + (TILE_N - 1)) / TILE_N;
+  if (GROUP_M > 1) {
+    grid_x = std::max(grid_x, static_cast<unsigned int>(GROUP_M));
+  }
+  bool DIVISIBLE_M = (M % TILE_M == 0);
+  bool DIVISIBLE_N = (N % TILE_N == 0);
+  bool DIVISIBLE_K = (K % TILE_K == 0);
+
+  f(/* CUstream = */ raw_stream,
+    /* grid_x = */ grid_x,
+    /* grid_y = */ grid_y,
+    /* grid_z = */ A_sizes[0],
+    /* num_warps = */ 4,
+    /* num_stages = */ 1,
+    A_contig,
+    B_contig,
+    out,
+    A_sizes[1],
+    B_sizes[2],
+    A_sizes[2],
+    TILE_M,
+    TILE_N,
+    TILE_K,
+    GROUP_M,
+    DIVISIBLE_M,
+    DIVISIBLE_N,
+    DIVISIBLE_K);
+  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("bmm(Tensor self, Tensor mat2) -> Tensor");
 }
 
 TORCH_LIBRARY_IMPL(flag_gems, CUDA, m) {
@@ -31,5 +32,6 @@ 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("bmm", TORCH_FN(bmm));
 }
 }  // namespace flag_gems

triton_src/bmm.py

@@ -0,0 +1,103 @@
+import triton
+import triton.language as tl
+
+from flag_gems.utils import triton_lang_extension as tle
+
+
+@triton.jit
+def bmm_kernel(
+    A,
+    B,
+    O,
+    M,
+    N,
+    K,
+    TILE_M: tl.constexpr,
+    TILE_N: tl.constexpr,
+    TILE_K: tl.constexpr,
+    GROUP_M: tl.constexpr,
+    DIVISIBLE_M: tl.constexpr,
+    DIVISIBLE_N: tl.constexpr,
+    DIVISIBLE_K: tl.constexpr,
+):
+    # batch offsets
+    pid_b = tle.program_id(2)
+    A += pid_b * M * K
+    B += pid_b * K * N
+    O += pid_b * M * N
+
+    pidx = tle.program_id(0)
+    pidy = tle.program_id(1)
+
+    if GROUP_M == 1:
+        pid_m, pid_n = pidx, pidy
+    else:
+        # reorder CTAs
+        gridx = tle.num_programs(0)
+        gridy = tle.num_programs(1)
+        pid = pidx + pidy * gridx
+
+        num_CTA_per_group = gridy * GROUP_M
+
+        group_id = pid // num_CTA_per_group
+        inner_group_id = pid % num_CTA_per_group
+        GROUP_SIZE = tl.where(
+            (group_id * GROUP_M + GROUP_M) > gridx, gridx % GROUP_M, GROUP_M
+        )
+        pid_m = group_id * GROUP_M + inner_group_id % GROUP_SIZE
+        pid_n = inner_group_id // GROUP_SIZE
+
+    offs_m = pid_m * TILE_M + tl.arange(0, TILE_M)
+    offs_n = pid_n * TILE_N + tl.arange(0, TILE_N)
+    offs_k = tl.arange(0, TILE_K)
+
+    if not DIVISIBLE_M:
+        mask_m = offs_m < M
+    if not DIVISIBLE_N:
+        mask_n = offs_n < N
+
+    a_ptrs = A + offs_m[:, None] * K + offs_k[None, :]
+    b_ptrs = B + offs_k[:, None] * N + offs_n[None, :]
+    o_ptrs = O + offs_m[:, None] * N + offs_n[None, :]
+
+    num_iters = tl.cdiv(K, TILE_K)
+    o = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
+    for _ in range(num_iters):
+        if DIVISIBLE_K:
+            if DIVISIBLE_M:
+                mask_a = None
+            else:
+                mask_a = mask_m[:, None]
+            if DIVISIBLE_N:
+                mask_b = None
+            else:
+                mask_b = mask_n[None, :]
+        else:
+            mask_k = offs_k < K
+            if DIVISIBLE_M:
+                mask_a = mask_k[None, :]
+            else:
+                mask_a = mask_m[:, None] & mask_k[None, :]
+            if DIVISIBLE_N:
+                mask_b = mask_k[:, None]
+            else:
+                mask_b = mask_k[:, None] & mask_n[None, :]
+
+        a = tl.load(a_ptrs, mask_a)
+        b = tl.load(b_ptrs, mask_b)
+
+        offs_k += TILE_K
+        a_ptrs += TILE_K
+        b_ptrs += TILE_K * N
+
+        o += tl.dot(a, b, allow_tf32=False)
+
+    if DIVISIBLE_M and DIVISIBLE_N:
+        mask_c = None
+    elif DIVISIBLE_M and not DIVISIBLE_N:
+        mask_c = mask_n[None, :]
+    elif not DIVISIBLE_M and DIVISIBLE_N:
+        mask_c = mask_m[:, None]
+    else:
+        mask_c = mask_m[:, None] & mask_n[None, :]
+    tl.store(o_ptrs, o, mask_c)