Generate ring masks directly as packed bits (#24089)

Summary:

Dense ring masks require an avoidable allocation and packing pass. Add `sdpa_bitwise_ring_mask_gen` schemas, meta/reference implementations and a Cadence C++ out kernel that writes packed `uint8` bits directly from token positions, sliding window and cache capacity. Keep the existing Python regressions and add seven public direct-call C++ unit tests for packed byte order and polarity, ring wrap, all windows on small caches, realistic cache sizes, `int32`/`int64` positions including `INT64_MAX`, empty inputs, invalid shapes/dtypes/windows/capacities and rejection of negative positions before writing. Wire the tests into the existing Buck and Cadence CMake test targets. This commit contains only ExecuTorch sources and synchronized mirrors; Jarvis registration remains in D126973553.

AI assistance: Codex separated the history and added the public C++ unit tests and build wiring. The production kernel is unchanged by this test addition.

Reviewed By: mvartani-meta

Differential Revision: D126734014

Pull Request resolved: #24089

Author: ThomasJannaud

backends/cadence/aot/ops_registrations.py

@@ -540,6 +540,13 @@ def register_fake(
     "sdpa_bitwise_causal_mask_gen.out(Tensor positions, int key_length, *, Tensor(a!) out) -> Tensor(a!)"
 )
 
+lib.define(
+    "sdpa_bitwise_ring_mask_gen(Tensor positions, int sliding_window, int cache_length) -> Tensor"
+)
+lib.define(
+    "sdpa_bitwise_ring_mask_gen.out(Tensor positions, int sliding_window, int cache_length, *, Tensor(a!) out) -> Tensor(a!)"
+)
+
 # Load/store with iDMA. These only exist before memory planning.
 # Post memory planning, we check that outputs/inputs for the load/store are in
 # DTCM and replace idma_load/idma_store with idma_copy.
@@ -3410,6 +3417,26 @@ def sdpa_bitwise_causal_mask_gen_meta(
     return positions.new_empty((positions.numel(), key_length // 8), dtype=torch.uint8)
 
 
+@register_fake("cadence::sdpa_bitwise_ring_mask_gen")
+def sdpa_bitwise_ring_mask_gen_meta(
+    positions: torch.Tensor, sliding_window: int, cache_length: int
+) -> torch.Tensor:
+    torch._check(
+        positions.dtype in (torch.int32, torch.int64),
+        lambda: "expected int32 or int64 positions",
+    )
+    torch._check(positions.dim() == 1, lambda: "positions must be one-dimensional")
+    torch._check(cache_length > 0, lambda: "cache_length must be positive")
+    torch._check(cache_length % 8 == 0, lambda: "cache_length must be a multiple of 8")
+    torch._check(
+        0 < sliding_window <= cache_length,
+        lambda: "sliding_window must be positive and at most cache_length",
+    )
+    return positions.new_empty(
+        (positions.numel(), cache_length // 8), dtype=torch.uint8
+    )
+
+
 @register_fake("cadence::quantized_w8a32_linear")
 def quantized_w8a32_linear_meta(
     src: torch.Tensor,

backends/cadence/aot/ref_implementations.py

@@ -3089,6 +3089,30 @@ def sdpa_bitwise_causal_mask_gen(
     ).to(torch.uint8)
 
 
+@impl_tracked(m, "sdpa_bitwise_ring_mask_gen")
+def sdpa_bitwise_ring_mask_gen(
+    positions: torch.Tensor, sliding_window: int, cache_length: int
+) -> torch.Tensor:
+    """Generate an LSB-first ring mask; set bits represent masked cache slots."""
+    assert positions.dtype in (torch.int32, torch.int64)
+    assert positions.dim() == 1
+    assert cache_length > 0 and cache_length % 8 == 0
+    assert 0 < sliding_window <= cache_length
+    assert bool(torch.all(positions >= 0))
+    absolute = positions.to(torch.int64).unsqueeze(1)
+    slots = torch.arange(
+        0, cache_length, 8, dtype=torch.int64, device=positions.device
+    ).unsqueeze(0)
+    packed = positions.new_zeros(
+        (positions.numel(), cache_length // 8), dtype=torch.uint8
+    )
+    for bit in range(8):
+        distance = (absolute - (slots + bit)) % cache_length
+        masked = (distance >= sliding_window) | (distance > absolute)
+        packed |= masked.to(torch.uint8) << bit
+    return packed
+
+
 @impl_tracked(m, "slice_scatter_")
 def slice_scatter_impl(
     self: torch.Tensor,

backends/cadence/aot/tests/test_ref_implementations.py

@@ -21,6 +21,67 @@
 
 
 class TestRefImplementations(unittest.TestCase):
+    def test_sdpa_bitwise_ring_mask_matches_visible_slots(self) -> None:
+        for dtype in (torch.int32, torch.int64):
+            for capacity in (8, 16, 24, 2048, 4096):
+                positions = torch.tensor(
+                    [0, 1, 7, 8, capacity - 1, capacity, 3 * capacity + 5],
+                    dtype=dtype,
+                )
+                for window in (1, min(9, capacity), capacity):
+                    with self.subTest(dtype=dtype, capacity=capacity, window=window):
+                        expected = torch.full(
+                            (positions.numel(), capacity // 8), 255, dtype=torch.uint8
+                        )
+                        for row, position in enumerate(positions.tolist()):
+                            for key in range(
+                                max(0, position - window + 1), position + 1
+                            ):
+                                slot = key % capacity
+                                expected[row, slot // 8] &= 255 ^ (1 << (slot % 8))
+                        actual = torch.ops.cadence.sdpa_bitwise_ring_mask_gen.default(
+                            positions, window, capacity
+                        )
+                        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+                        meta = torch.ops.cadence.sdpa_bitwise_ring_mask_gen.default(
+                            positions.to("meta"), window, capacity
+                        )
+                        self.assertEqual(meta.shape, expected.shape)
+                        self.assertEqual(meta.dtype, torch.uint8)
+
+    def test_sdpa_bitwise_ring_mask_empty_and_large_positions(self) -> None:
+        empty = torch.ops.cadence.sdpa_bitwise_ring_mask_gen.default(
+            torch.empty(0, dtype=torch.int64), 4, 16
+        )
+        self.assertEqual(tuple(empty.shape), (0, 2))
+        actual = torch.ops.cadence.sdpa_bitwise_ring_mask_gen.default(
+            torch.tensor([torch.iinfo(torch.int64).max]), 4, 16
+        )
+        torch.testing.assert_close(
+            actual, torch.tensor([[255, 15]], dtype=torch.uint8), rtol=0, atol=0
+        )
+
+    def test_sdpa_bitwise_ring_mask_rejects_invalid_inputs(self) -> None:
+        for device in ("cpu", "meta"):
+            cases = [
+                (torch.zeros(1, dtype=torch.float32, device=device), 4, 16),
+                (torch.zeros((1, 1), dtype=torch.int64, device=device), 4, 16),
+                (torch.zeros(1, dtype=torch.int64, device=device), 0, 16),
+                (torch.zeros(1, dtype=torch.int64, device=device), 17, 16),
+                (torch.zeros(1, dtype=torch.int64, device=device), 4, 0),
+                (torch.zeros(1, dtype=torch.int64, device=device), 4, 15),
+            ]
+            for positions, window, capacity in cases:
+                with self.subTest(device=device, window=window, capacity=capacity):
+                    with self.assertRaises((AssertionError, RuntimeError)):
+                        torch.ops.cadence.sdpa_bitwise_ring_mask_gen.default(
+                            positions, window, capacity
+                        )
+        with self.assertRaises(AssertionError):
+            torch.ops.cadence.sdpa_bitwise_ring_mask_gen.default(
+                torch.tensor([-1], dtype=torch.int64), 4, 16
+            )
+
     @expand(
         [
             ("basic_int8", 0.42, -1.0, 2.0, -7, 7, torch.int8, 0),

backends/cadence/generic/operators/op_sdpa_bitwise_ring_mask_gen.cpp

@@ -0,0 +1,122 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/cadence/generic/operators/op_sdpa_bitwise_ring_mask_gen.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+
+#include <executorch/runtime/kernel/kernel_includes.h>
+
+namespace impl {
+namespace generic {
+namespace native {
+namespace {
+
+void clear_range(uint8_t* out, int64_t start, int64_t end) {
+  if (start == end) {
+    return;
+  }
+  const int64_t first_byte = start / 8;
+  const int64_t last_byte = (end - 1) / 8;
+  if (first_byte == last_byte) {
+    const uint32_t visible = ((1u << (end - start)) - 1) << (start % 8);
+    out[first_byte] &= static_cast<uint8_t>(~visible);
+    return;
+  }
+  out[first_byte] &= static_cast<uint8_t>((1u << (start % 8)) - 1);
+  std::memset(
+      out + first_byte + 1,
+      0,
+      static_cast<std::size_t>(last_byte - first_byte - 1));
+  out[last_byte] &= static_cast<uint8_t>(0xFFu << ((end - 1) % 8 + 1));
+}
+
+template <typename Position>
+bool generate_mask(
+    const Position* positions,
+    int64_t num_positions,
+    int64_t sliding_window,
+    int64_t cache_length,
+    uint8_t* out) {
+  for (int64_t row = 0; row < num_positions; ++row) {
+    if (positions[row] < 0) {
+      return false;
+    }
+  }
+  const int64_t packed_width = cache_length / 8;
+  for (int64_t row = 0; row < num_positions; ++row) {
+    uint8_t* row_out = out + row * packed_width;
+    std::memset(row_out, 0xFF, static_cast<std::size_t>(packed_width));
+    const int64_t position = positions[row];
+    const int64_t visible =
+        position >= sliding_window - 1 ? sliding_window : position + 1;
+    const int64_t end = position % cache_length + 1;
+    const int64_t start = end - visible;
+    if (start < 0) {
+      clear_range(row_out, cache_length + start, cache_length);
+      clear_range(row_out, 0, end);
+    } else {
+      clear_range(row_out, start, end);
+    }
+  }
+  return true;
+}
+
+} // namespace
+
+::executorch::aten::Tensor& sdpa_bitwise_ring_mask_gen_out(
+    ::executorch::runtime::KernelRuntimeContext& ctx,
+    const ::executorch::aten::Tensor& positions,
+    int64_t sliding_window,
+    int64_t cache_length,
+    ::executorch::aten::Tensor& out) {
+  using ::executorch::aten::ScalarType;
+  ET_KERNEL_CHECK(ctx, positions.dim() == 1, InvalidArgument, out);
+  ET_KERNEL_CHECK(
+      ctx, cache_length > 0 && cache_length % 8 == 0, InvalidArgument, out);
+  ET_KERNEL_CHECK(
+      ctx,
+      sliding_window > 0 && sliding_window <= cache_length,
+      InvalidArgument,
+      out);
+  ET_KERNEL_CHECK(
+      ctx,
+      positions.dtype() == ScalarType::Int ||
+          positions.dtype() == ScalarType::Long,
+      InvalidArgument,
+      out);
+  ET_KERNEL_CHECK(ctx, out.dtype() == ScalarType::Byte, InvalidArgument, out);
+  ET_KERNEL_CHECK(
+      ctx,
+      out.dim() == 2 && out.size(0) == positions.numel() &&
+          out.size(1) == cache_length / 8,
+      InvalidArgument,
+      out);
+
+  const bool valid = positions.dtype() == ScalarType::Long
+      ? generate_mask(
+            positions.const_data_ptr<int64_t>(),
+            positions.numel(),
+            sliding_window,
+            cache_length,
+            out.mutable_data_ptr<uint8_t>())
+      : generate_mask(
+            positions.const_data_ptr<int32_t>(),
+            positions.numel(),
+            sliding_window,
+            cache_length,
+            out.mutable_data_ptr<uint8_t>());
+  ET_KERNEL_CHECK(ctx, valid, InvalidArgument, out);
+  return out;
+}
+
+} // namespace native
+} // namespace generic
+} // namespace impl

backends/cadence/generic/operators/op_sdpa_bitwise_ring_mask_gen.h

@@ -0,0 +1,29 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/kernel/kernel_runtime_context.h>
+
+namespace impl {
+namespace generic {
+namespace native {
+
+::executorch::aten::Tensor& sdpa_bitwise_ring_mask_gen_out(
+    ::executorch::runtime::KernelRuntimeContext& ctx,
+    const ::executorch::aten::Tensor& positions,
+    int64_t sliding_window,
+    int64_t cache_length,
+    ::executorch::aten::Tensor& out);
+
+} // namespace native
+} // namespace generic
+} // namespace impl

backends/cadence/generic/operators/targets.bzl

@@ -520,6 +520,19 @@ def define_common_targets():
         visibility = ["PUBLIC"],
     )
 
+    runtime.cxx_library(
+        name = "op_sdpa_bitwise_ring_mask_gen",
+        srcs = ["op_sdpa_bitwise_ring_mask_gen.cpp"],
+        exported_headers = ["op_sdpa_bitwise_ring_mask_gen.h"],
+        platforms = CXX,
+        deps = [
+            "//executorch/runtime/kernel:kernel_includes",
+            "//executorch/runtime/core/exec_aten:lib",
+            "//executorch/runtime/kernel:kernel_runtime_context",
+        ],
+        visibility = ["PUBLIC"],
+    )
+
     runtime.cxx_library(
         name = "op_sdpa_bitwise_mask_gen",
         srcs = ["op_sdpa_bitwise_mask_gen.cpp"],

backends/cadence/generic/operators/tests/BUCK

@@ -41,3 +41,22 @@ jarvis_wrapper.cxx_test(
         "fbsource//xplat/executorch/runtime/core/exec_aten/testing_util:tensor_util",
     ],
 )
+
+jarvis_wrapper.cxx_test(
+    name = "test_op_sdpa_bitwise_ring_mask_gen",
+    srcs = [
+        "test_op_sdpa_bitwise_ring_mask_gen.cpp",
+    ],
+    compatible_backends = ["generic"],
+    labels = [ci.linux(ci.mode("fbsource//arvr/mode/platform010/dev"))],
+    platforms = CXX,
+    visibility = [
+        "fbsource//xplat/executorch/backends/cadence/...",
+        "fbcode//executorch/backends/cadence/...",
+    ],
+    deps = [
+        "fbsource//xplat/executorch/backends/cadence/generic/operators:op_sdpa_bitwise_ring_mask_gen",
+        "fbsource//xplat/executorch/kernels/test:gtest_utils",
+        "fbsource//xplat/executorch/runtime/core/exec_aten/testing_util:tensor_util",
+    ],
+)

backends/cadence/generic/operators/tests/CMakeLists.txt

@@ -49,6 +49,7 @@ endif()
 add_executable(
   cadence_generic_op_tests
   test_op_quantized_layer_norm.cpp test_op_sdpa_bitwise_mask_gen.cpp
+  test_op_sdpa_bitwise_ring_mask_gen.cpp
   ${EXECUTORCH_ROOT}/runtime/core/exec_aten/testing_util/tensor_util.cpp
 )
 target_compile_definitions(cadence_generic_op_tests PRIVATE GTEST_HAS_PTHREAD=0)