Implement IsInf JIT Emitter for ARM64 SIMD in OpenVINO (#24471)

Hello maintainers,

I've implemented the IsInf JIT emitter for the ARM64 SIMD platform, as
outlined in the OpenVINO CPU plugin JIT emitters documentation.

This PR addresses
[#24419](https://github.com/openvinotoolkit/openvino/issues/24419).

Please let me know if any further adjustments are needed.
Thank you!

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_eltwise_emitters.cpp

@@ -801,6 +801,105 @@ std::set<std::vector<element::Type>> jit_hswish_emitter::get_supported_precision
     return {{element::f32}};
 }
 
+/// IS_INF ///
+
+jit_is_inf_emitter::jit_is_inf_emitter(dnnl::impl::cpu::aarch64::jit_generator* host,
+                                       dnnl::impl::cpu::aarch64::cpu_isa_t host_isa,
+                                       const std::shared_ptr<ov::Node>& node)
+    : jit_emitter(host, host_isa, node, get_arithmetic_binary_exec_precision(node)) {
+
+    auto isInf = ov::as_type_ptr<ov::op::v10::IsInf>(node);
+    if (isInf == nullptr) {
+        OV_CPU_JIT_EMITTER_THROW("Can't cast to ov::op::v10::IsInf");
+    }
+
+    const auto& attributes = isInf->get_attributes();
+    detect_negative = attributes.detect_negative;
+    detect_positive = attributes.detect_positive;
+
+    prepare_table();
+}
+
+jit_is_inf_emitter::jit_is_inf_emitter(dnnl::impl::cpu::aarch64::jit_generator* host,
+                                       dnnl::impl::cpu::aarch64::cpu_isa_t host_isa,
+                                       const bool detect_negative,
+                                       const bool detect_positive,
+                                       const ov::element::Type exec_prc)
+    : jit_emitter(host, host_isa, exec_prc),
+      detect_negative{detect_negative},
+      detect_positive{detect_positive} {
+    prepare_table();
+}
+
+size_t jit_is_inf_emitter::get_inputs_count() const {
+    return 1;
+}
+
+size_t jit_is_inf_emitter::get_aux_vecs_count() const {
+    return 1;
+}
+
+size_t jit_is_inf_emitter::get_aux_gprs_count() const {
+    return 1;
+}
+
+std::set<std::vector<element::Type>> jit_is_inf_emitter::get_supported_precisions(
+    const std::shared_ptr<ov::Node>& node) {
+    return {{element::f32}};
+}
+
+void jit_is_inf_emitter::emit_impl(const std::vector<size_t>& in_vec_idxs,
+                                   const std::vector<size_t>& out_vec_idxs) const {
+    if (host_isa_ == dnnl::impl::cpu::aarch64::asimd) {
+        emit_isa<dnnl::impl::cpu::aarch64::asimd>(in_vec_idxs, out_vec_idxs);
+    } else {
+        OV_CPU_JIT_EMITTER_THROW("Can't create jit eltwise kernel");
+    }
+}
+
+template <dnnl::impl::cpu::aarch64::cpu_isa_t isa>
+void jit_is_inf_emitter::emit_isa(const std::vector<size_t>& in_vec_idxs,
+                                  const std::vector<size_t>& out_vec_idxs) const {
+    OV_CPU_JIT_EMITTER_ASSERT(exec_prc_ == ov::element::f32, "unsupported precision: " + exec_prc_.to_string());
+
+    using TReg = typename dnnl::impl::cpu::aarch64::cpu_isa_traits<isa>::TReg;
+    const TReg src = TReg(in_vec_idxs[0]);
+    const TReg dst = TReg(out_vec_idxs[0]);
+    const TReg aux = TReg(aux_vec_idxs[0]);
+
+    if (detect_negative || detect_positive) {
+        if (detect_positive) {
+            if (detect_negative) {
+                // If both positive and negative infinity detection is requested
+                // calculate the absolute value of 'src'.
+                h->fabs(src.s, src.s);
+            }
+            // Load 'aux' with positive infinity.
+            h->ld1r(aux.s, table_val2("inf"));
+        } else if (detect_negative) {
+            // Load 'aux' with negative infinity.
+            h->ld1r(aux.s, table_val2("inf_neg"));
+        }
+        // Compare elements of 'src' with 'aux'.
+        h->fcmeq(dst.s, src.s, aux.s);
+        // Sets elements in 'dst' to 1.0 where the comparison was true.
+        h->ld1r(aux.s, table_val2("one"));
+        h->and_(dst.b16, dst.b16, aux.b16);
+
+    } else {
+        // If neither positive nor negative infinity detection is enabled,
+        // set 'dst' with zeros (a eor a is 0)
+        h->eor(dst.b16, dst.b16, dst.b16);
+    }
+}
+
+void jit_is_inf_emitter::register_table_entries() {
+    // Registers constant values that comply with the IEEE 754 standard.
+    push_arg_entry_of("one", 0x3F800000, true);
+    push_arg_entry_of("inf", 0x7F800000, true);
+    push_arg_entry_of("inf_neg", 0xFF800000, true);
+}
+
 /// MAX ///
 jit_maximum_emitter::jit_maximum_emitter(dnnl::impl::cpu::aarch64::jit_generator* host,
                                          dnnl::impl::cpu::aarch64::cpu_isa_t host_isa,

src/plugins/intel_cpu/src/emitters/plugin/aarch64/jit_eltwise_emitters.hpp

@@ -378,6 +378,39 @@ private:
     void emit_isa(const std::vector<size_t> &in_vec_idxs, const std::vector<size_t> &out_vec_idxs) const;
 };
 
+class jit_is_inf_emitter : public jit_emitter {
+public:
+    jit_is_inf_emitter(dnnl::impl::cpu::aarch64::jit_generator* host,
+                       dnnl::impl::cpu::aarch64::cpu_isa_t host_isa,
+                       const std::shared_ptr<ov::Node>& node);
+
+    jit_is_inf_emitter(dnnl::impl::cpu::aarch64::jit_generator* host,
+                       dnnl::impl::cpu::aarch64::cpu_isa_t host_isa,
+                       const bool detect_negative,
+                       const bool detect_positive,
+                       const ov::element::Type exec_prc = ov::element::f32);
+
+    size_t get_inputs_count() const override;
+
+    size_t get_aux_vecs_count() const override;
+
+    size_t get_aux_gprs_count() const override;
+
+    static std::set<std::vector<element::Type>> get_supported_precisions(
+        const std::shared_ptr<ov::Node>& node = nullptr);
+
+private:
+    void emit_impl(const std::vector<size_t>& in_vec_idxs, const std::vector<size_t>& out_vec_idxs) const override;
+
+    template <dnnl::impl::cpu::aarch64::cpu_isa_t isa>
+    void emit_isa(const std::vector<size_t>& in_vec_idxs, const std::vector<size_t>& out_vec_idxs) const;
+
+    void register_table_entries() override;
+
+    bool detect_negative;
+    bool detect_positive;
+};
+
 class jit_mod_emitter : public jit_emitter {
 public:
     jit_mod_emitter(dnnl::impl::cpu::aarch64::jit_generator *host,

src/plugins/intel_cpu/src/nodes/executors/aarch64/jit_eltwise.cpp

@@ -29,6 +29,7 @@ bool JitEltwiseExecutor::isSupported(
                                      Algorithm::EltwiseGeluErf,
                                      Algorithm::EltwiseGeluTanh,
                                      Algorithm::EltwiseHswish,
+                                     Algorithm::EltwiseIsInf,
                                      Algorithm::EltwiseMaximum,
                                      Algorithm::EltwiseMinimum,
                                      Algorithm::EltwiseMish,

src/plugins/intel_cpu/src/nodes/kernels/aarch64/jit_uni_eltwise_generic.cpp

@@ -618,6 +618,17 @@ struct EltwiseEmitter<jit_power_static_emitter> {
     }
 };
 
+template<>
+struct EltwiseEmitter<jit_is_inf_emitter> {
+    void operator()(EltwiseEmitterContext& ctx) {
+        ctx.emitter = std::make_shared<jit_is_inf_emitter>(ctx.host,
+                                                           ctx.host_isa,
+                                                           ctx.opData.alpha,
+                                                           ctx.opData.beta,
+                                                           ctx.exec_prc);
+    }
+};
+
 template <dnnl::impl::cpu::aarch64::cpu_isa_t isa>
 std::shared_ptr<jit_emitter> jit_uni_eltwise_generic<isa>::create_eltwise_emitter(const EltwiseData& data, const ov::element::Type& exec_prec) {
     EltwiseEmitterContext ctx = {
@@ -638,6 +649,7 @@ std::shared_ptr<jit_emitter> jit_uni_eltwise_generic<isa>::create_eltwise_emitte
     OV_CASE(Algorithm::EltwiseExp, ov::intel_cpu::aarch64::jit_exp_emitter),
     OV_CASE(Algorithm::EltwiseFloor, ov::intel_cpu::aarch64::jit_floor_emitter),
     OV_CASE(Algorithm::EltwiseHswish, ov::intel_cpu::aarch64::jit_hswish_emitter),
+    OV_CASE(Algorithm::EltwiseIsInf, ov::intel_cpu::aarch64::jit_is_inf_emitter),
     OV_CASE(Algorithm::EltwiseMaximum, ov::intel_cpu::aarch64::jit_maximum_emitter),
     OV_CASE(Algorithm::EltwiseMinimum, ov::intel_cpu::aarch64::jit_minimum_emitter),
     OV_CASE(Algorithm::EltwiseMish, ov::intel_cpu::aarch64::jit_mish_emitter),
@@ -810,6 +822,7 @@ std::set<std::vector<element::Type>> eltwise_precision_helper::get_supported_pre
         OV_CASE(Algorithm::EltwiseGeluErf, jit_gelu_erf_emitter),
         OV_CASE(Algorithm::EltwiseGeluTanh, jit_gelu_tanh_emitter),
         OV_CASE(Algorithm::EltwiseHswish, jit_hswish_emitter),
+        OV_CASE(Algorithm::EltwiseIsInf, jit_is_inf_emitter),
         OV_CASE(Algorithm::EltwiseMaximum, jit_maximum_emitter),
         OV_CASE(Algorithm::EltwiseMinimum, jit_minimum_emitter),
         OV_CASE(Algorithm::EltwiseMish, jit_mish_emitter),

src/plugins/intel_cpu/tests/functional/custom/single_layer_tests/classes/activation.cpp

@@ -161,6 +161,7 @@ std::string ActivationLayerCPUTest::getPrimitiveType(const utils::ActivationType
         (activation_type == utils::ActivationTypes::Exp) ||
         (activation_type == utils::ActivationTypes::Floor) ||
         (activation_type == utils::ActivationTypes::HSwish) ||
+        (activation_type == utils::ActivationTypes::IsInf) ||
         (activation_type == utils::ActivationTypes::HardSigmoid) ||
         (activation_type == utils::ActivationTypes::Mish) ||
         (activation_type == utils::ActivationTypes::GeluErf) ||

src/plugins/intel_cpu/tests/functional/shared_tests_instances/single_layer_tests/activation.cpp

@@ -54,7 +54,8 @@ const std::map<ActivationTypes, std::vector<std::vector<float>>> activationTypes
         {ActivationTypes::RoundHalfAwayFromZero, {}},
         {ActivationTypes::GeluErf,               {}},
         {ActivationTypes::GeluTanh,              {}},
-        {ActivationTypes::Swish,                 {{0.4f}}}
+        {ActivationTypes::Swish,                 {{0.4f}}},
+        {ActivationTypes::IsInf,                 {}}
 };
 
 // List of operations that should be tested also with integer precision

src/tests/functional/shared_test_classes/include/shared_test_classes/single_op/activation.hpp

@@ -67,6 +67,7 @@ static std::map<ActivationTypes, std::string> activationNames = {
         {ActivationTypes::GeluErf,               "GeluErf"},
         {ActivationTypes::GeluTanh,              "GeluTanh"},
         {ActivationTypes::SoftSign,              "SoftSign"},
+        {ActivationTypes::IsInf,                 "IsInf"},
 };
 
 typedef std::tuple<

src/tests/test_utils/common_test_utils/include/common_test_utils/test_enums.hpp

@@ -121,7 +121,8 @@ enum ActivationTypes {
     RoundHalfAwayFromZero,
     GeluErf,
     GeluTanh,
-    SoftSign
+    SoftSign,
+    IsInf
 };
 
 enum MinMaxOpType {

src/tests/test_utils/common_test_utils/src/node_builders/activation.cpp

@@ -24,6 +24,7 @@
 #include "openvino/op/hard_sigmoid.hpp"
 #include "openvino/op/hsigmoid.hpp"
 #include "openvino/op/hswish.hpp"
+#include "openvino/op/is_inf.hpp"
 #include "openvino/op/log.hpp"
 #include "openvino/op/mish.hpp"
 #include "openvino/op/negative.hpp"
@@ -144,6 +145,8 @@ std::shared_ptr<ov::Node> make_activation(const ov::Output<Node>& in,
         return std::make_shared<ov::op::v7::Gelu>(in, ov::op::GeluApproximationMode::TANH);
     case ov::test::utils::ActivationTypes::SoftSign:
         return std::make_shared<ov::op::v9::SoftSign>(in);
+    case ov::test::utils::ActivationTypes::IsInf:
+        return std::make_shared<ov::op::v10::IsInf>(in);
     default:
         OPENVINO_THROW("Can't create layer for this activation type");
     }