feat: support distinct bool tensor dtype in InfiniOps

docs/api/core-types.md

@@ -46,6 +46,7 @@ surface.
 - `DataType::kInt8`, `DataType::kInt16`, `DataType::kInt32`, `DataType::kInt64`
 - `DataType::kUInt8`, `DataType::kUInt16`, `DataType::kUInt32`, `DataType::kUInt64`
 - `DataType::kFloat16`, `DataType::kBFloat16`, `DataType::kFloat32`, `DataType::kFloat64`
+- `DataType::kBool` (distinct from `kUInt8`, with one byte per element)
 
 InfiniOps also exposes type-list helpers such as `FloatTypes`,
 `ReducedFloatTypes`, `IntTypes`, `UIntTypes`, and `AllTypes` for template

src/base/topk_sigmoid.h

@@ -162,7 +162,7 @@ class TopkSigmoid : public Operator<TopkSigmoid> {
     if (is_padding) {
       assert(is_padding->ndim() == 1 && is_padding->numel() == num_tokens_ &&
              "`TopkSigmoid` requires `is_padding` shape `[num_tokens]`");
-      assert(is_padding->dtype() == DataType::kUInt8 &&
+      assert(is_padding->dtype() == DataType::kBool &&
              is_padding->IsContiguous() &&
              "`TopkSigmoid` requires contiguous bool `is_padding`");
       assert(same_device(*is_padding) &&

src/base/topk_softmax.h

@@ -145,7 +145,7 @@ class TopkSoftmax : public Operator<TopkSoftmax> {
     if (is_padding) {
       assert(is_padding->ndim() == 1 && is_padding->numel() == num_tokens_ &&
              "`TopkSoftmax` requires `is_padding` shape `[num_tokens]`");
-      assert(is_padding->dtype() == DataType::kUInt8 &&
+      assert(is_padding->dtype() == DataType::kBool &&
              is_padding->IsContiguous() &&
              "`TopkSoftmax` requires contiguous bool `is_padding`");
       assert(same_device(*is_padding) &&

src/pybind11_utils.h

@@ -104,10 +104,6 @@ std::unordered_map<std::string, Device::Type> BuildTorchNameMap(
 }  // namespace detail
 
 inline DataType DataTypeFromString(std::string_view name) {
-  // InfiniRT has no bool dtype; carry bool tensor storage as byte data and
-  // restore bool semantics in operators that accept bool tensors.
-  if (name == "bool") return DataType::kUInt8;
-
   return kStringToDataType.at(name);
 }
 

src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc

@@ -39,12 +39,9 @@ void Operator<ScaledDotProductAttention, kDev, 1>::operator()(
 
   c10::optional<at::Tensor> at_attn_mask;
   if (attn_mask.has_value()) {
-    const auto dtype_override = attn_mask_type_ == DataType::kUInt8
-                                    ? std::optional<at::ScalarType>{at::kBool}
-                                    : std::nullopt;
     at_attn_mask.emplace(ToAtenTensor<kDev>(
         const_cast<void*>(attn_mask->data()), attn_mask_shape_,
-        attn_mask_strides_, attn_mask_type_, device_index_, dtype_override));
+        attn_mask_strides_, attn_mask_type_, device_index_));
   }
 
   c10::optional<double> at_scale;

src/torch/tensor_.h

@@ -63,6 +63,8 @@ inline at::ScalarType ToAtenDataType(DataType dtype) {
       return at::kLong;
     case DataType::kUInt8:
       return at::kByte;
+    case DataType::kBool:
+      return at::kBool;
     case DataType::kUInt16:
     case DataType::kUInt32:
     case DataType::kUInt64:

tests/test_all.py

@@ -0,0 +1,15 @@
+import infini.ops
+import pytest
+import torch
+
+
+@pytest.mark.parametrize("input_dtype", (torch.bool, torch.uint8))
+def test_bool_and_uint8_reductions_remain_distinct(input_dtype, device):
+    input = torch.tensor([1, 0, 2], dtype=input_dtype, device=device)
+    expected = torch.all(input, dim=0, keepdim=False)
+    out = torch.empty_like(expected)
+
+    infini.ops.all(input, 0, False, out, implementation_index=8)
+
+    assert out.dtype == expected.dtype
+    torch.testing.assert_close(out, expected, rtol=0, atol=0)

tests/test_any.py

@@ -0,0 +1,15 @@
+import infini.ops
+import pytest
+import torch
+
+
+@pytest.mark.parametrize("input_dtype", (torch.bool, torch.uint8))
+def test_bool_and_uint8_reductions_remain_distinct(input_dtype, device):
+    input = torch.tensor([1, 0, 2], dtype=input_dtype, device=device)
+    expected = torch.any(input, dim=0, keepdim=False)
+    out = torch.empty_like(expected)
+
+    infini.ops.any(input, 0, False, out, implementation_index=8)
+
+    assert out.dtype == expected.dtype
+    torch.testing.assert_close(out, expected, rtol=0, atol=0)

tests/test_isin.py

@@ -0,0 +1,16 @@
+import infini.ops
+import pytest
+import torch
+
+
+@pytest.mark.parametrize("input_dtype", (torch.uint8, torch.float32))
+def test_bool_output(input_dtype, device):
+    elements = torch.tensor([0, 1, 2, 1], dtype=input_dtype, device=device)
+    test_elements = torch.tensor([1, 2], dtype=input_dtype, device=device)
+    expected = torch.isin(elements, test_elements, assume_unique=False, invert=False)
+    out = torch.empty_like(expected)
+
+    infini.ops.isin(elements, test_elements, False, False, out, implementation_index=8)
+
+    assert out.dtype == torch.bool
+    torch.testing.assert_close(out, expected, rtol=0, atol=0)

tests/test_logical_not.py

@@ -0,0 +1,17 @@
+import infini.ops
+import pytest
+import torch
+
+
+@pytest.mark.parametrize("input_dtype", (torch.bool, torch.uint8))
+@pytest.mark.parametrize("strided", (False, True))
+def test_bool_output_distinct_from_uint8(input_dtype, strided, device):
+    values = torch.tensor([[0, 1, 2], [1, 0, 0]], dtype=input_dtype, device=device)
+    input = values.t() if strided else values
+    expected = torch.logical_not(input)
+    out = torch.empty_like(expected)
+
+    infini.ops.logical_not(input, out, implementation_index=8)
+
+    assert out.dtype == torch.bool
+    torch.testing.assert_close(out, expected, rtol=0, atol=0)

tests/test_torch_ops.py

@@ -114,7 +114,7 @@
 _TYPE_DEFAULTS = {"int": 0, "SymInt": 0, "bool": False, "str": "none"}
 
 # Mirrors `kStringToDataType` in `src/data_type.h`.  Any tensor passed to
-# an InfiniOps op must have one of these dtypes; others (`bool`, complex,
+# an InfiniOps op must have one of these dtypes; others (complex,
 # quantised types) abort the process inside `DataTypeFromString`.  Some
 # vendor torch forks lag behind upstream and lack `uint16` / `uint32` /
 # `uint64` (added in PyTorch 2.3); resolve them lazily and keep the
@@ -128,6 +128,7 @@
     "uint16",
     "uint32",
     "uint64",
+    "bool",
     "float16",
     "bfloat16",
     "float32",
@@ -679,8 +680,8 @@ def test_op(op_meta, shape, dtype, device, rtol, atol):
         )
 
     # InfiniOps `DataType` supports only `int{8,16,32,64}`,
-    # `uint{8,16,32,64}`, `float{16,32,64}`, and `bfloat16`.  Tensors with
-    # any other torch dtype (`bool`, `complex64`, `complex128`, etc.) abort
+    # `uint{8,16,32,64}`, `bool`, `float{16,32,64}`, and `bfloat16`. Tensors with
+    # any other torch dtype (`complex64`, `complex128`, etc.) abort
     # on `DataTypeFromString`, so skip the test rather than crash the process.
     tensors = [*ref_outs, *(x for x in inputs if isinstance(x, torch.Tensor))]
     unsupported = next(