diff --git a/docs/api/core-types.md b/docs/api/core-types.md index ab24946cc..a9cb6919e 100644 --- a/docs/api/core-types.md +++ b/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 diff --git a/src/base/topk_sigmoid.h b/src/base/topk_sigmoid.h index 76a8f59b0..2c569c4c8 100644 --- a/src/base/topk_sigmoid.h +++ b/src/base/topk_sigmoid.h @@ -162,7 +162,7 @@ class TopkSigmoid : public Operator { 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) && diff --git a/src/base/topk_softmax.h b/src/base/topk_softmax.h index 92722c09e..43aaac734 100644 --- a/src/base/topk_softmax.h +++ b/src/base/topk_softmax.h @@ -145,7 +145,7 @@ class TopkSoftmax : public Operator { 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) && diff --git a/src/pybind11_utils.h b/src/pybind11_utils.h index 1bd6f0bcd..4acc97d37 100644 --- a/src/pybind11_utils.h +++ b/src/pybind11_utils.h @@ -104,10 +104,6 @@ std::unordered_map 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); } diff --git a/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc b/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc index db5cf10f4..543ca6da0 100644 --- a/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc +++ b/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc @@ -39,12 +39,9 @@ void Operator::operator()( c10::optional at_attn_mask; if (attn_mask.has_value()) { - const auto dtype_override = attn_mask_type_ == DataType::kUInt8 - ? std::optional{at::kBool} - : std::nullopt; at_attn_mask.emplace(ToAtenTensor( const_cast(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 at_scale; diff --git a/src/torch/tensor_.h b/src/torch/tensor_.h index 556eed5cd..a1cea20fd 100644 --- a/src/torch/tensor_.h +++ b/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: diff --git a/tests/test_all.py b/tests/test_all.py new file mode 100644 index 000000000..3cf13a4f5 --- /dev/null +++ b/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) diff --git a/tests/test_any.py b/tests/test_any.py new file mode 100644 index 000000000..adce77c1d --- /dev/null +++ b/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) diff --git a/tests/test_isin.py b/tests/test_isin.py new file mode 100644 index 000000000..ab1be0f2c --- /dev/null +++ b/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) diff --git a/tests/test_logical_not.py b/tests/test_logical_not.py new file mode 100644 index 000000000..8bd28a52e --- /dev/null +++ b/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) diff --git a/tests/test_torch_ops.py b/tests/test_torch_ops.py index c5576ca3c..dcf045a86 100644 --- a/tests/test_torch_ops.py +++ b/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(