[AdvancedCompiler]Zeros(cpp wrapper)

CMakeLists.txt

@@ -54,6 +54,7 @@ if(FLAGGEMS_USE_EXTERNAL_TRITON_JIT)
 else()
   set(TRITON_JIT_INSTALL ON) # install triton jit
   FetchContent_Declare(TritonJIT
+    #GIT_REPOSITORY  git@github.com:AdvancedCompiler/libtriton_jit
     GIT_REPOSITORY https://github.com/iclementine/libtorch_example.git
     # SOURCE_DIR /home/clement/projects/libtorch_example    # use local source dir in development
   )

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_zeros test_triton_zeros.cpp)
+target_link_libraries(test_triton_zeros
+  PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)
+add_test(NAME test_triton_zeros COMMAND test_triton_zeros)

ctests/test_triton_reduction.cpp

@@ -2,7 +2,6 @@
 #include "c10/util/Logging.h"
 #include "flag_gems/operators.h"
 #include "torch/torch.h"
-
 TEST(reduction_op_test, sum) {
   const torch::Device device(torch::kCUDA, 0);
   torch::Tensor a = torch::randn({32, 1024}, device);

ctests/test_triton_zeros.cpp

@@ -0,0 +1,17 @@
+#include <gtest/gtest.h>
+
+#include "flag_gems/operators.h"
+#include "torch/torch.h"
+TEST(zeros_op_test, zeros) {
+  const torch::Device device(torch::kCUDA, 0);
+  int64_t n = 100;
+  auto options = torch::TensorOptions().device(device).dtype(torch::kFloat32);
+  torch::Tensor out_torch = torch::zeros({n}, options);
+  torch::Tensor out_triton_0 = flag_gems::zeros_tensor(n);
+  torch::Tensor out_triton_1 = flag_gems::zeros_tensor(n, torch::kFloat32);
+  torch::Tensor out_triton_2 = flag_gems::zeros_tensor(n, torch::kFloat32, c10::nullopt);
+
+  EXPECT_TRUE(torch::allclose(out_torch, out_triton_0));
+  EXPECT_TRUE(torch::allclose(out_torch, out_triton_1));
+  EXPECT_TRUE(torch::allclose(out_torch, out_triton_2));
+}

include/flag_gems/operators.h

@@ -3,6 +3,9 @@
 #include "torch/torch.h"
 
 namespace flag_gems {
+at::Tensor zeros_tensor(int64_t n_elements,
+                        ::std::optional<at::ScalarType> dtype = ::std::nullopt,
+                        ::std::optional<at::Device> device = ::std::nullopt);
 at::Tensor add_tensor(const at::Tensor &a_, const at::Tensor &b_);
 at::Tensor mm_tensor(const at::Tensor &mat1, const at::Tensor &mat2);
 at::Tensor sum_dim(const at::Tensor &self,

lib/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_library(operators
             SHARED
+	          zeros.cpp
             utils.cpp
             add.cpp
             sum.cpp

lib/zeros.cpp

@@ -0,0 +1,43 @@
+#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 zeros_tensor(int64_t n_elements,
+                        c10::optional<at::ScalarType> dtype,
+                        c10::optional<at::Device> device) {
+  TORCH_CHECK(n_elements > 0, "the element of tensor must >0")
+
+  at::ScalarType final_dtype = dtype.value_or(at::typeMetaToScalarType(at::get_default_dtype()));
+  at::Device final_device =
+      device.value_or(torch::cuda::is_available() ? at::Device(at::kCUDA) : at::Device(at::kCPU));
+  at::Tensor out = at::empty({n_elements}, at::TensorOptions().dtype(final_dtype).device(final_device));
+  int64_t tile_size = 1024;
+  const int num_warps = 8;
+  const int num_stages = 1;
+
+  const uint64_t num_blocks = (static_cast<uint64_t>(n_elements) + tile_size - 1) / tile_size;
+  const TritonJITFunction &f =
+      TritonJITFunction::getInstance(std::string(utils::get_triton_src_path() / "zeros.py"), "zeros_kernel");
+
+  c10::DeviceGuard guard(out.device());
+  c10::cuda::CUDAStream stream = c10::cuda::getCurrentCUDAStream();
+  CUstream raw_stream = static_cast<CUstream>(stream.stream());
+
+  f(stream,
+    num_blocks,
+    /* grid_y = */ 1,
+    /* grid_z = */ 1,
+    /* num_warps = */ num_warps,
+    /* num_stages = */ num_stages,
+    out,
+    n_elements,
+    tile_size);
+  return out;
+}
+}  // namespace flag_gems

src/flag_gems/csrc/cstub.cpp

@@ -8,6 +8,7 @@ PYBIND11_MODULE(c_operators, m) {
 
 namespace flag_gems {
 TORCH_LIBRARY(flag_gems, m) {
+  m.def("zeros_tensor(SymInt n_elements, ScalarType? dtype=None, Device? device=None) -> Tensor");
   m.def("sum.dim_IntList(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor");
   m.def("add_tensor(Tensor self, Tensor other) -> Tensor", {at::Tag::pt2_compliant_tag});
   // Norm
@@ -23,6 +24,7 @@ TORCH_LIBRARY(flag_gems, m) {
 }
 
 TORCH_LIBRARY_IMPL(flag_gems, CUDA, m) {
+  m.impl("zeros_tensor", TORCH_FN(zeros_tensor));
   m.impl("sum.dim_IntList", TORCH_FN(sum_dim));
   m.impl("add_tensor", TORCH_FN(add_tensor));
   // Norm

triton_src/zeros.py

@@ -0,0 +1,17 @@
+import triton
+import triton.language as tl
+
+from flag_gems.utils import triton_lang_extension as tle
+
+
+@triton.jit
+def zeros_kernel(
+    output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tle.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    tl.store(output_ptr + offsets, 0.0, mask=mask)