Add Int4 and UInt4 support for Cast

onnxruntime/core/providers/cpu/tensor/cast_op.cc

@@ -262,6 +262,58 @@
   }
 };
 
+// tensor Int4x2 -> float
+template <>
+struct TensorCaster<Int4x2, float> {
+  void Cast(const OpKernelContext& ctx, const TensorShape& shape, const Tensor& in, Tensor& out) const {
+    const auto* in_data = in.Data<Int4x2>();
+    auto* out_data = out.MutableData<float>();
+
+    // Confirm we can unpack the int4
+    const size_t shape_size = narrow<size_t>(shape.Size());
+    const size_t in_shape_size = narrow<size_t>(in.Shape().Size());
+    ORT_ENFORCE(in_shape_size * 2 == shape_size,
+                "The Int4x2 tensor size is invalid for casting to float.");
+
+    for (size_t i = 0; i < in_shape_size; ++i) {
+      const uint8_t packed = static_cast<uint8_t>(in_data[i].bits_);
+
+      // Extract signed high and low nibble
+      int8_t high_nibble = static_cast<int8_t>(packed) >> 4;
+      int8_t low_nibble = static_cast<int8_t>(packed << 4) >> 4;
+
+      out_data[2 * i] = static_cast<float>(high_nibble);
+      out_data[2 * i + 1] = static_cast<float>(low_nibble);
+    }
+  }
+};
+
+// tensor UInt4x2 -> float
+template <>
+struct TensorCaster<UInt4x2, float> {
+  void Cast(const OpKernelContext& ctx, const TensorShape& shape, const Tensor& in, Tensor& out) const {
+    const auto* in_data = in.Data<UInt4x2>();
+    auto* out_data = out.MutableData<float>();
+
+    // Confirm we can unpack the uint4
+    const size_t shape_size = narrow<size_t>(shape.Size());
+    const size_t in_shape_size = narrow<size_t>(in.Shape().Size());
+    ORT_ENFORCE(in_shape_size * 2 == shape_size,
+                "The UInt4x2 tensor size is invalid for casting to float.");
+
+    for (size_t i = 0; i < in_shape_size; ++i) {
+      const uint8_t packed = static_cast<uint8_t>(in_data[i].bits_);
+
+      // Extract unsigned high and low nibble
+      uint8_t high_nibble = (packed >> 4) & 0x0F;
+      uint8_t low_nibble = packed & 0x0F;
+
+      out_data[2 * i] = static_cast<float>(high_nibble);
+      out_data[2 * i + 1] = static_cast<float>(low_nibble);
+    }
+  }
+};
+
 #if defined(_M_AMD64) && !defined(_M_ARM64EC)
 // specializations to use optimized and Windows x64-specific
 

onnxruntime/test/providers/cpu/tensor/cast_op_test.cc

@@ -207,6 +207,38 @@ TEST(CastOpTest, ToString) {
   TestCastOp(gsl::make_span(int_16_input), gsl::make_span(int_string_data), shape);
 }
 
+TEST(CastOpTest, Int4ToFloat) {
+  // GIVEN
+  const std::vector<int64_t> shape{2, 2, 2};
+  const std::vector<Int4x2> int4_input = {
+      Int4x2(1, 2), // two 4-bit int elements: lower = 1, upper = 2
+      Int4x2(-3, -4),
+      Int4x2(5, -6),
+      Int4x2(-8, 7)
+  };
+  // There will be twice as many unpacked elements
+  const std::vector<float> expected_float_output = {1.0f, 2.0f, -3.0f, -4.0f, 5.0f, -6.0f, -8.0f, 7.0f};
+
+  // WHEN, THEN
+  TestCastOp(gsl::make_span(int4_input), gsl::make_span(expected_float_output), shape);
+}
+
+TEST(CastOpTest, UInt4ToFloat) {
+  // GIVEN
+  const std::vector<int64_t> shape{2, 2, 2};
+  const std::vector<UInt4x2> uint4_input = {
+      UInt4x2(0, 1),
+      UInt4x2(2, 3),
+      UInt4x2(7, 8),
+      UInt4x2(14, 15)
+  };
+  // There will be twice as many unpacked elements
+  const std::vector<float> expected_float_output = {0.0f, 1.0f, 2.0f, 3.0f, 7.0f, 8.0f, 14.0f, 15.0f};
+
+  // WHEN, THEN
+  TestCastOp(gsl::make_span(uint4_input), gsl::make_span(expected_float_output), shape);
+}
+
 #if !defined(DISABLE_FLOAT8_TYPES)
 
 template <typename F8>