[Bug] Fixed instruction RW info related to AVX-512 {k} (merging) and {k} {z} (zeroing)

2 files changed, 48 insertions, 12 deletions

diff --git a/src/asmjit/x86/x86instapi.cpp b/src/asmjit/x86/x86instapi.cpp
index fec69b2..114bca2 100644
--- a/src/asmjit/x86/x86instapi.cpp
+++ b/src/asmjit/x86/x86instapi.cpp
@@ -809,6 +809,20 @@ static ASMJIT_INLINE void rwZeroExtendNonVec(OpRWInfo& opRwInfo, const Reg& reg)
   }
 }
 
+static ASMJIT_INLINE Error rwHandleAVX512(const BaseInst& inst, InstRWInfo* out) noexcept {
+  if (inst.hasExtraReg() && inst.extraReg().type() == Reg::kTypeKReg && out->opCount() > 0) {
+    // AVX-512 instruction that uses a destination with {k} register (zeroing vs masking).
+    out->_extraReg.addOpFlags(OpRWInfo::kRead);
+    out->_extraReg.setReadByteMask(0xFF);
+    if (!inst.hasOption(Inst::kOptionZMask)) {
+      out->_operands[0].addOpFlags(OpRWInfo::kRead);
+      out->_operands[0]._readByteMask |= out->_operands[0]._writeByteMask;
+    }
+  }
+
+  return kErrorOk;
+}
+
 Error InstInternal::queryRWInfo(uint32_t arch, const BaseInst& inst, const Operand_* operands, size_t opCount, InstRWInfo* out) noexcept {
   using namespace Status;
 
@@ -940,15 +954,15 @@ Error InstInternal::queryRWInfo(uint32_t arch, const BaseInst& inst, const Opera
       } while (it.hasNext());
     }
 
-    return kErrorOk;
+    return rwHandleAVX512(inst, out);
   }
 
   switch (instRwInfo.category) {
     case InstDB::RWInfo::kCategoryMov: {
-      // Special case for 'movhpd' instruction. Here there are some variants that
-      // we have to handle as mov can be used to move between GP, segment, control
-      // and debug registers. Moving between GP registers also allow to use memory
-      // operand.
+      // Special case for 'mov' instruction. Here there are some variants that
+      // we have to handle as 'mov' can be used to move between GP, segment,
+      // control and debug registers. Moving between GP registers also allow to
+      // use memory operand.
 
       if (opCount == 2) {
         if (operands[0].isReg() && operands[1].isReg()) {
@@ -1118,7 +1132,7 @@ Error InstInternal::queryRWInfo(uint32_t arch, const BaseInst& inst, const Opera
     case InstDB::RWInfo::kCategoryMovh64: {
       // Special case for 'movhpd|movhps' instructions. Note that this is only
       // required for legacy (non-AVX) variants as AVX instructions use either
-      // 2 or 3 operands that are use `kCategoryGeneric`.
+      // 2 or 3 operands that are in `kCategoryGeneric` category.
       if (opCount == 2) {
         if (BaseReg::isVec(operands[0]) && operands[1].isMem()) {
           out->_operands[0].reset(W, 8);
@@ -1174,7 +1188,7 @@ Error InstInternal::queryRWInfo(uint32_t arch, const BaseInst& inst, const Opera
           out->_operands[1]._readByteMask &= 0x00FF00FF00FF00FFu;
 
           rwZeroExtendAvxVec(out->_operands[0], operands[0].as<Vec>());
-          return kErrorOk;
+          return rwHandleAVX512(inst, out);
         }
 
         if (BaseReg::isVec(operands[0]) && operands[1].isMem()) {
@@ -1185,7 +1199,7 @@ Error InstInternal::queryRWInfo(uint32_t arch, const BaseInst& inst, const Opera
           out->_operands[1].reset(R | MibRead, o1Size);
 
           rwZeroExtendAvxVec(out->_operands[0], operands[0].as<Vec>());
-          return kErrorOk;
+          return rwHandleAVX512(inst, out);
         }
       }
       break;
@@ -1259,7 +1273,7 @@ Error InstInternal::queryRWInfo(uint32_t arch, const BaseInst& inst, const Opera
           if (BaseReg::isVec(operands[0]))
             rwZeroExtendAvxVec(out->_operands[0], operands[0].as<Vec>());
 
-          return kErrorOk;
+          return rwHandleAVX512(inst, out);
         }
 
         if (operands[0].isReg() && operands[1].isMem()) {
@@ -1277,7 +1291,8 @@ Error InstInternal::queryRWInfo(uint32_t arch, const BaseInst& inst, const Opera
 
           out->_operands[0].reset(W | MibRead, size0);
           out->_operands[1].reset(R, size1);
-          return kErrorOk;
+
+          return rwHandleAVX512(inst, out);
         }
       }
       break;
@@ -1328,12 +1343,14 @@ Error InstInternal::queryRWInfo(uint32_t arch, const BaseInst& inst, const Opera
             out->_operands[1].addOpFlags(RegM);
             out->_operands[1].setRmSize(size1);
           }
-          return kErrorOk;
+
+          return rwHandleAVX512(inst, out);
         }
 
         if (operands[0].isReg() && operands[1].isMem()) {
           out->_operands[1].addOpFlags(MibRead);
-          return kErrorOk;
+
+          return rwHandleAVX512(inst, out);
         }
       }
       break;
diff --git a/test/asmjit_test_x86_instinfo.cpp b/test/asmjit_test_x86_instinfo.cpp
index 8f58263..68d870d 100644
--- a/test/asmjit_test_x86_instinfo.cpp
+++ b/test/asmjit_test_x86_instinfo.cpp
@@ -148,6 +148,15 @@ static void printInfoSimple(uint32_t arch, uint32_t instId, Args&&... args) {
   printInfo(arch, inst, opArray, sizeof...(args));
 }
 
+template<typename... Args>
+static void printInfoExtra(uint32_t arch, uint32_t instId, uint32_t options, const BaseReg& extraReg, Args&&... args) {
+  BaseInst inst(instId);
+  inst.addOptions(options);
+  inst.setExtraReg(extraReg);
+  Operand_ opArray[] = { std::forward<Args>(args)... };
+  printInfo(arch, inst, opArray, sizeof...(args));
+}
+
 static void testX86Arch() {
 #if defined(ASMJIT_BUILD_X86)
   uint32_t arch = Environment::kArchX64;
@@ -183,6 +192,16 @@ static void testX86Arch() {
   printInfoSimple(arch,
                   x86::Inst::kIdVaddpd,
                   x86::zmm0, x86::zmm1, x86::zmm2);
+
+  printInfoExtra(arch,
+                 x86::Inst::kIdVaddpd, 0,
+                 x86::k1,
+                 x86::zmm0, x86::zmm1, x86::zmm2);
+
+  printInfoExtra(arch,
+                 x86::Inst::kIdVaddpd, x86::Inst::kOptionZMask,
+                 x86::k1,
+                 x86::zmm0, x86::zmm1, x86::zmm2);
 #endif
 }