Skip to content

Add signed BitShift support for opset 28 #32361

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mirounga merged 2 commits into from
Sep 21, 2026
Merged

Add signed BitShift support for opset 28 #32361

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ec9b0ed
Add signed BitShift support for opset 28
justinchuby Sep 1, 2026
a1635b2
Address BitShift opset 28 review feedback
justinchuby Sep 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 40 additions & 9 deletions onnxruntime/core/providers/cpu/cpu_execution_provider.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -596,9 +596,10 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOn
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, MLFloat16, Gemm);
#endif
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, GatherElements);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, uint8_t, BitShift);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, uint32_t, BitShift);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, uint64_t, BitShift);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 27, uint8_t, BitShift);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 27, uint16_t, BitShift);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 27, uint32_t, BitShift);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 27, uint64_t, BitShift);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, Pad);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 11, GatherND);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 26, Range);
Expand Down Expand Up @@ -1549,6 +1550,16 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain,
// Opset 27
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 27, Range);

// Opset 28
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, int8_t, BitShift);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, int16_t, BitShift);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, int32_t, BitShift);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, int64_t, BitShift);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, uint8_t, BitShift);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, uint16_t, BitShift);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, uint32_t, BitShift);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, uint64_t, BitShift);

// !!PLEASE READ BELOW!! Following that, add new entries above this comment

/* *** IMPORTANT! ***
Expand Down Expand Up @@ -2311,12 +2322,14 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
double, Gemm)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12,
GatherElements)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, uint8_t,
BitShift)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, uint32_t,
BitShift)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, uint64_t,
BitShift)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 27,
uint8_t, BitShift)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 27,
uint16_t, BitShift)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 27,
uint32_t, BitShift)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 27,
uint64_t, BitShift)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, Pad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 11,
GatherND)>,
Expand Down Expand Up @@ -3742,6 +3755,24 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {

// opset 27
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 27, Range)>,

// opset 28
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, int8_t,
BitShift)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, int16_t,
BitShift)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, int32_t,
BitShift)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, int64_t,
BitShift)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, uint8_t,
BitShift)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, uint16_t,
BitShift)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, uint32_t,
BitShift)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 28, uint64_t,
BitShift)>,
};
for (auto& function_table_entry : function_table) {
KernelCreateInfo info = function_table_entry();
Expand Down
66 changes: 57 additions & 9 deletions onnxruntime/core/providers/cpu/math/element_wise_ops.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@
#include "core/mlas/inc/mlas.h"

#include <cmath>
#include <cstring>

Check warning on line 15 in onnxruntime/core/providers/cpu/math/element_wise_ops.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 Found C++ system header after other header. Should be: element_wise_ops.h, c system, c++ system, other.  [build/include_order] [4]

Raw Output:
onnxruntime/core/providers/cpu/math/element_wise_ops.cc:15:  Found C++ system header after other header. Should be: element_wise_ops.h, c system, c++ system, other.  [build/include_order] [4]
#include <limits>

Check warning on line 16 in onnxruntime/core/providers/cpu/math/element_wise_ops.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 Found C++ system header after other header. Should be: element_wise_ops.h, c system, c++ system, other.  [build/include_order] [4]

Raw Output:
onnxruntime/core/providers/cpu/math/element_wise_ops.cc:16:  Found C++ system header after other header. Should be: element_wise_ops.h, c system, c++ system, other.  [build/include_order] [4]
#include <type_traits>

Check warning on line 17 in onnxruntime/core/providers/cpu/math/element_wise_ops.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 Found C++ system header after other header. Should be: element_wise_ops.h, c system, c++ system, other.  [build/include_order] [4]

Raw Output:
onnxruntime/core/providers/cpu/math/element_wise_ops.cc:17:  Found C++ system header after other header. Should be: element_wise_ops.h, c system, c++ system, other.  [build/include_order] [4]

namespace onnxruntime {
// Supported types for operators that have type reduction enabled
Expand Down Expand Up @@ -489,10 +492,18 @@
// Supposed to add BFloat16 but we are not supporting now, however, separate registration
REG_ELEMENTWISE_TYPED_KERNEL(Mean, 13, float, Mean_8);

REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 11, uint8_t, BitShift);
// REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 11, uint16_t, BitShift);
REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 11, uint32_t, BitShift);
REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 11, uint64_t, BitShift);
REG_ELEMENTWISE_VERSIONED_TYPED_KERNEL(BitShift, 11, 27, uint8_t, BitShift);
REG_ELEMENTWISE_VERSIONED_TYPED_KERNEL(BitShift, 11, 27, uint16_t, BitShift);
REG_ELEMENTWISE_VERSIONED_TYPED_KERNEL(BitShift, 11, 27, uint32_t, BitShift);
REG_ELEMENTWISE_VERSIONED_TYPED_KERNEL(BitShift, 11, 27, uint64_t, BitShift);
REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, int8_t, BitShift);
Comment thread
justinchuby marked this conversation as resolved.
REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, int16_t, BitShift);
REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, int32_t, BitShift);
REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, int64_t, BitShift);
REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, uint8_t, BitShift);
REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, uint16_t, BitShift);
REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, uint32_t, BitShift);
REG_ELEMENTWISE_TYPED_KERNEL(BitShift, 28, uint64_t, BitShift);

REG_ELEMENTWISE_TYPED_KERNEL(BitwiseAnd, 18, int8_t, BitwiseAnd);
REG_ELEMENTWISE_TYPED_KERNEL(BitwiseAnd, 18, int16_t, BitwiseAnd);
Expand Down Expand Up @@ -1334,17 +1345,54 @@
ORT_THROW("Invalid direction value of '", direction, "'. Valid values are 'LEFT' or 'RIGHT'.");
}

// Shifting by >= the bit width of an unsigned type is undefined behavior in C++.
// On x86, 64-bit shifts mask the shift amount to 6 bits, so shift by 64 acts like shift by 0.
// Guard against this by returning 0 when the shift amount >= the bit width.
template <typename T>
T FromUnsignedBits(std::make_unsigned_t<T> bits) {
static_assert(std::is_integral_v<T>);
T value;
std::memcpy(&value, &bits, sizeof(value));
return value;
}

template <typename T>
inline bool IsInvalidShift(T shift) {
using UnsignedT = std::make_unsigned_t<T>;
constexpr UnsignedT bit_width = std::numeric_limits<UnsignedT>::digits;
return (std::is_signed_v<T> && shift < 0) || static_cast<UnsignedT>(shift) >= bit_width;
}

template <typename T>
inline T SafeShiftLeft(T value, T shift) {
return shift >= sizeof(T) * 8 ? T{0} : value << shift;
if (IsInvalidShift(shift)) {
return T{0};
}

using UnsignedT = std::make_unsigned_t<T>;
const auto result = static_cast<UnsignedT>(value) << static_cast<UnsignedT>(shift);
if constexpr (std::is_signed_v<T>) {
return FromUnsignedBits<T>(result);
} else {
return result;
}
Comment thread
justinchuby marked this conversation as resolved.
Outdated
}

template <typename T>
inline T SafeShiftRight(T value, T shift) {
return shift >= sizeof(T) * 8 ? T{0} : value >> shift;
if (IsInvalidShift(shift)) {
return std::is_signed_v<T> && value < 0 ? T{-1} : T{0};
}

using UnsignedT = std::make_unsigned_t<T>;
constexpr UnsignedT bit_width = std::numeric_limits<UnsignedT>::digits;
const auto shift_amount = static_cast<UnsignedT>(shift);
auto result = static_cast<UnsignedT>(value) >> shift_amount;
if constexpr (std::is_signed_v<T>) {
if (value < 0 && shift_amount != 0) {
result |= std::numeric_limits<UnsignedT>::max() << (bit_width - shift_amount);
}
return FromUnsignedBits<T>(result);
} else {
return result;
}
}

template <typename T>
Expand Down
5 changes: 5 additions & 0 deletions onnxruntime/core/providers/dml/DmlExecutionProvider/src/AbiCustomRegistry.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -391,6 +391,11 @@ HRESULT STDMETHODCALLTYPE AbiCustomRegistry::RegisterOperatorKernel(
builder.SetDomain(opKernel->domain)
.SinceVersion(opKernel->minimumOperatorSetVersion)
.Provider(providerType);
if (std::string_view(opKernel->name) == "BitShift" && opKernel->minimumOperatorSetVersion == 11)
{
// DirectML masks out-of-range shift counts, which does not satisfy the opset-28 semantics.
builder.SinceVersion(11, 27);
Comment thread
justinchuby marked this conversation as resolved.
}

std::string_view name(opKernel->name);
if (name == "MemcpyToHost")
Expand Down
2 changes: 0 additions & 2 deletions onnxruntime/test/onnx/TestCase.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1091,8 +1091,6 @@ std::unique_ptr<std::set<BrokenTest>> GetBrokenTests(const std::string& provider
{"resize_upsample_sizes_nearest_ceil_half_pixel", "Bad onnx test output. Needs test fix."},
{"resize_upsample_sizes_nearest_floor_align_corners", "Bad onnx test output. Needs test fix."},
{"resize_upsample_sizes_nearest_round_prefer_ceil_asymmetric", "Bad onnx test output. Needs test fix."},
Comment thread
justinchuby marked this conversation as resolved.
{"bitshift_right_uint16", "BitShift(11) uint16 support not enabled currently"},
{"bitshift_left_uint16", "BitShift(11) uint16 support not enabled currently"},
{"maxunpool_export_with_output_shape",
"Invalid output in ONNX test. See https://github.com/onnx/onnx/issues/2398"},
{"cntk_simple_seg", "Bad onnx test output caused by wrong SAME_UPPER/SAME_LOWER for ConvTranspose"},
Expand Down
79 changes: 79 additions & 0 deletions onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
#include <algorithm>
#include <limits>
#include <math.h>
#include <type_traits>

#if defined(USE_WEBGPU)
// Pull in the deviceless shared-trailing-dimension helper from the WebGPU EP so it can be
Expand Down Expand Up @@ -4921,6 +4922,84 @@ TEST(BitShiftOpTest, BroadcastXRight_Uint8) {
test.Run();
}

template <typename T>
void TestSignedBitShiftOpset28() {
constexpr T bit_width = static_cast<T>(std::numeric_limits<std::make_unsigned_t<T>>::digits);

{
OpTester test("BitShift", 28);
test.AddAttribute("direction", "LEFT");
test.AddInput<T>("X", {6}, {T{-8}, T{-3}, T{1}, std::numeric_limits<T>::max(), T{-2}, T{2}});
test.AddInput<T>("Y", {6}, {T{1}, T{2}, T{3}, T{1}, T{-1}, bit_width});
test.AddOutput<T>("Z", {6}, {T{-16}, T{-12}, T{8}, T{-2}, T{0}, T{0}});
test.Run();
}

{
OpTester test("BitShift", 28);
test.AddAttribute("direction", "RIGHT");
test.AddInput<T>("X", {6}, {T{-8}, T{-3}, T{1}, T{-2}, T{-2}, T{2}});
test.AddInput<T>("Y", {6}, {T{1}, T{2}, T{3}, T{-1}, bit_width, bit_width});
test.AddOutput<T>("Z", {6}, {T{-4}, T{-1}, T{0}, T{-1}, T{-1}, T{0}});
test.Run();
}

{
OpTester test("BitShift", 28);
test.AddAttribute("direction", "RIGHT");
test.AddInput<T>("X", {1}, {T{-8}});
test.AddInput<T>("Y", {4}, {T{1}, T{2}, bit_width, T{-1}});
test.AddOutput<T>("Z", {4}, {T{-4}, T{-2}, T{-1}, T{-1}});
test.Run();
}

{
OpTester test("BitShift", 28);
test.AddAttribute("direction", "LEFT");
test.AddInput<T>("X", {2}, {std::numeric_limits<T>::max(), T{-8}});
test.AddInput<T>("Y", {1}, {T{1}});
test.AddOutput<T>("Z", {2}, {T{-2}, T{-16}});
test.Run();
}
}

TEST(BitShiftOpTest, SignedInt8Opset28) {
TestSignedBitShiftOpset28<int8_t>();
}

TEST(BitShiftOpTest, SignedInt16Opset28) {
TestSignedBitShiftOpset28<int16_t>();
}

TEST(BitShiftOpTest, SignedInt32Opset28) {
TestSignedBitShiftOpset28<int32_t>();
}

TEST(BitShiftOpTest, SignedInt64Opset28) {
TestSignedBitShiftOpset28<int64_t>();
}

template <typename T>
void TestUnsignedBitShiftRegistration(int opset) {
OpTester test("BitShift", opset);
test.AddAttribute("direction", "LEFT");
test.AddInput<T>("X", {2}, {T{1}, T{3}});
test.AddInput<T>("Y", {2}, {T{1}, T{2}});
test.AddOutput<T>("Z", {2}, {T{2}, T{12}});
test.Run();
}

TEST(BitShiftOpTest, Uint16Opset11) {
TestUnsignedBitShiftRegistration<uint16_t>(11);
}

TEST(BitShiftOpTest, UnsignedTypesOpset28) {
TestUnsignedBitShiftRegistration<uint8_t>(28);
TestUnsignedBitShiftRegistration<uint16_t>(28);
TestUnsignedBitShiftRegistration<uint32_t>(28);
TestUnsignedBitShiftRegistration<uint64_t>(28);
}
Comment thread
justinchuby marked this conversation as resolved.

// Test that shift amounts >= bit width produce 0 (not undefined behavior).
// DirectML EP has the same hardware-level shift masking behavior, so skip these tests for DML.
TEST(BitShiftOpTest, RightShiftByBitWidth_Uint64) {
Expand Down
Loading