[AdvancedCompiler]Zeros(cpp wrapper)

ctests/CMakeLists.txt

@@ -26,6 +26,11 @@ 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_zeros_constructor test_triton_tensor_constructor.cpp)
+target_link_libraries(test_triton_zeros_constructor
+  PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)
+add_test(NAME test_triton_zeros_constructor COMMAND test_triton_zeros_constructor)
 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,44 @@
+#include <gtest/gtest.h>
+
+#include "flag_gems/operators.h"
+#include "torch/torch.h"
+
+TEST(zeros_op_test, 2d_tensor) {
+  const torch::Device device(torch::kCUDA, 0);
+  std::vector<int64_t> shape_0 = {31};
+  std::vector<int64_t> shape_1 = {11, 7};
+  std::vector<int64_t> shape = {7, 7, 7};
+  auto options = torch::TensorOptions().device(device).dtype(torch::kFloat32);
+
+  torch::Tensor ref_empty = torch::empty(shape, options);
+  torch::Tensor ref_empty_0 = torch::empty(shape_0, options);
+  torch::Tensor ref_empty_1 = torch::empty(shape_1, options);
+  ref_empty.fill_(0);
+  ref_empty_0.fill_(0);
+  ref_empty_1.fill_(0);
+  torch::Tensor out_triton = flag_gems::zeros(torch::IntArrayRef(shape),  // size
+                                              torch::kFloat32,            // dtype
+                                              c10::nullopt,               // layout
+                                              device                      // device
+  );
+
+  torch::Tensor out_triton_0 = flag_gems::zeros(torch::IntArrayRef(shape_0),  // size
+                                                torch::kFloat32,              // dtype
+                                                c10::nullopt,                 // layout
+                                                device                        // device
+  );
+  torch::Tensor out_triton_1 = flag_gems::zeros(torch::IntArrayRef(shape_1),  // size
+                                                torch::kFloat32,              // dtype
+                                                c10::nullopt,                 // layout
+                                                device                        // device
+  );
+
+  EXPECT_TRUE(torch::all(out_triton == 0).item<bool>());
+  EXPECT_TRUE(torch::allclose(out_triton, ref_empty));
+
+  EXPECT_TRUE(torch::all(out_triton_0 == 0).item<bool>());
+  EXPECT_TRUE(torch::allclose(out_triton_0, ref_empty_0));
+
+  EXPECT_TRUE(torch::all(out_triton_1 == 0).item<bool>());
+  EXPECT_TRUE(torch::allclose(out_triton_1, ref_empty_1));
+}

include/flag_gems/operators.h

@@ -3,6 +3,11 @@
 #include "torch/torch.h"
 
 namespace flag_gems {
+at::Tensor zeros(at::IntArrayRef size,
+                 c10::optional<at::ScalarType> dtype = ::std::nullopt,
+                 c10::optional<at::Layout> layout = ::std::nullopt,
+                 c10::optional<at::Device> device = ::std::nullopt,
+                 c10::optional<bool> pin_memory = ::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,61 @@
+#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(at::IntArrayRef size,
+                 c10::optional<at::ScalarType> dtype,
+                 c10::optional<at::Layout> layout,
+                 c10::optional<at::Device> device,
+                 c10::optional<bool> pin_memory) {
+  int64_t n_elements = 1;
+  for (auto dim : size) {
+    n_elements *= dim;
+  }
+
+  auto options =
+      at::TensorOptions()
+          .dtype(dtype.value_or(at::typeMetaToScalarType(at::get_default_dtype())))
+          .layout(layout.value_or(at::kStrided))
+          .device(device.value_or(torch::cuda::is_available() ? at::Device(at::kCUDA) : at::Device(at::kCPU)))
+          .pinned_memory(pin_memory.value_or(false));
+
+  TORCH_CHECK(n_elements >= 0, "Total elements must be non-negative");
+
+  if (n_elements == 0) {
+    return at::empty(size, options);
+  }
+
+  at::Tensor out = at::empty(size, options);
+
+  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

@@ -19,6 +19,9 @@ PYBIND11_MODULE(c_operators, m) {
 
 namespace flag_gems {
 TORCH_LIBRARY(flag_gems, m) {
+  m.def(
+      "zeros(SymInt[] size, ScalarType? dtype=None,Layout? layout=None, Device? device=None, bool? "
+      "pin_memory=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
@@ -37,6 +40,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)