Added push(...) and pop(...) for std::vector<> of SIMD registers

Author: lohika-denis-kotov

src/cpu/x64/jit_generator.hpp

@@ -200,6 +200,34 @@ class jit_generator : public Xbyak::CodeGenerator, public c_compatible {
         }
     }
 
+    inline void push(const std::vector<Xbyak::Xmm> &xmms) {
+        std::vector<std::function<void()>> deferred_movs{};
+        size_t offset = 0;
+        for (size_t i = 0; i < xmms.size(); ++i) {
+            const auto& xmm = xmms[i];
+            if (xmm.isXMM()) {
+                deferred_movs.emplace_back([this, offset, &xmm]() {
+                    uni_vmovdqu(ptr[rsp + offset], xmm);
+                });
+                offset += xmm_len;
+            } else if (xmm.isYMM()) {
+                deferred_movs.emplace_back([this, offset, &xmm]() {
+                    uni_vmovdqu(ptr[rsp + offset], Xbyak::Ymm{xmm.getIdx()});
+                });
+                offset += ymm_len;
+            } else if (xmm.isZMM()) {
+                deferred_movs.emplace_back([this, offset, &xmm]() {
+                    uni_vmovdqu(ptr[rsp + offset], Xbyak::Zmm{xmm.getIdx()});
+                });
+                offset += zmm_len;
+            }
+        }
+        sub(rsp, offset);
+        for (const auto& def_mov : deferred_movs) {
+            def_mov();
+        }
+    }
+
     inline void pop(const Xbyak::Xmm &xmm) {
         if (xmm.isXMM()) {
             uni_vmovdqu(xmm, ptr[rsp]);
@@ -213,6 +241,34 @@ class jit_generator : public Xbyak::CodeGenerator, public c_compatible {
         }
     }
 
+    inline void pop(const std::vector<Xbyak::Xmm> &xmms) {
+        std::vector<std::function<void()>> deferred_movs{};
+        size_t offset = 0;
+        for (size_t i = 0; i < xmms.size(); ++i) {
+            const auto& xmm = xmms[i];
+            if (xmm.isXMM()) {
+                deferred_movs.emplace_back([this, offset, &xmm]() {
+                    uni_vmovdqu(xmm, ptr[rsp + offset]);
+                });
+                offset += xmm_len;
+            } else if (xmm.isYMM()) {
+                deferred_movs.emplace_back([this, offset, &xmm]() {
+                    uni_vmovdqu(Xbyak::Ymm{xmm.getIdx()}, ptr[rsp + offset]);
+                });
+                offset += ymm_len;
+            } else if (xmm.isZMM()) {
+                deferred_movs.emplace_back([this, offset, &xmm]() {
+                    uni_vmovdqu(Xbyak::Zmm{xmm.getIdx()}, ptr[rsp + offset]);
+                });
+                offset += zmm_len;
+            }
+        }
+        for (const auto& def_mov : deferred_movs) {
+            def_mov();
+        }
+        add(rsp, offset);
+    }
+
     void preamble() {
         if (xmm_to_preserve) {
             sub(rsp, xmm_to_preserve * xmm_len);