[AdvancedCompiler]Zeros(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_zeros test_triton_tensor_constructor.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)
 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_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_tensor_constructor.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(n);
+  torch::Tensor out_triton_1 = flag_gems::zeros(n, torch::kFloat32);
+  torch::Tensor out_triton_2 = flag_gems::zeros(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(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,41 @@
+#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(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(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
@@ -24,6 +25,7 @@ TORCH_LIBRARY(flag_gems, m) {
 }
 
 TORCH_LIBRARY_IMPL(flag_gems, CUDA, m) {
+  m.impl("zeros", TORCH_FN(zeros));
   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)