[X86][MC]Fix wrong action for encode movdir64b

Movdir64b is special for its mem operand, 67 prefex can not only modify its add size,
so it's mem base and index reg should be the same type as source reg, such as
movdir64b (%rdx), rcx, and could not be movdir64b (%edx), rcx.
Now llvm-mc can encode the asm 'movdir64b (%edx), rcx' but the result is the same as
'movdir64b (%edx), ecx', which offend users' intention, while gcc will object this
action and give a warning.
I add 3 new mem descriptions to let llvm-mc to report the same error.

Reviewed By: skan, craig.topper

Differential Revision: https://reviews.llvm.org/D145893

5 files changed, 77 insertions, 3 deletions

diff --git a/llvm/lib/Target/X86/AsmParser/X86Operand.h b/llvm/lib/Target/X86/AsmParser/X86Operand.h
index 075b800f9e20..4661e73c3ef8 100644
--- a/llvm/lib/Target/X86/AsmParser/X86Operand.h
+++ b/llvm/lib/Target/X86/AsmParser/X86Operand.h
@@ -380,6 +380,40 @@ struct X86Operand final : public MCParsedAsmOperand {
   bool isMem512_RC512() const {
     return isMem512() && isMemIndexReg(X86::ZMM0, X86::ZMM31);
   }
+  bool isMem512_GR16() const {
+    if (!isMem512())
+      return false;
+    if (getMemBaseReg() &&
+        !X86MCRegisterClasses[X86::GR16RegClassID].contains(getMemBaseReg()))
+      return false;
+    return true;
+  }
+  bool isMem512_GR32() const {
+    if (!isMem512())
+      return false;
+    if (getMemBaseReg() &&
+        !X86MCRegisterClasses[X86::GR32RegClassID].contains(getMemBaseReg()) &&
+        getMemBaseReg() != X86::EIP)
+      return false;
+    if (getMemIndexReg() &&
+        !X86MCRegisterClasses[X86::GR32RegClassID].contains(getMemIndexReg()) &&
+        getMemIndexReg() != X86::EIZ)
+      return false;
+    return true;
+  }
+  bool isMem512_GR64() const {
+    if (!isMem512())
+      return false;
+    if (getMemBaseReg() &&
+        !X86MCRegisterClasses[X86::GR64RegClassID].contains(getMemBaseReg()) &&
+        getMemBaseReg() != X86::RIP)
+      return false;
+    if (getMemIndexReg() &&
+        !X86MCRegisterClasses[X86::GR64RegClassID].contains(getMemIndexReg()) &&
+        getMemIndexReg() != X86::RIZ)
+      return false;
+    return true;
+  }
 
   bool isAbsMem() const {
     return Kind == Memory && !getMemSegReg() && !getMemBaseReg() &&
diff --git a/llvm/lib/Target/X86/X86InstrInfo.td b/llvm/lib/Target/X86/X86InstrInfo.td
index b832b15c5b93..d295208c75f7 100644
--- a/llvm/lib/Target/X86/X86InstrInfo.td
+++ b/llvm/lib/Target/X86/X86InstrInfo.td
@@ -380,6 +380,9 @@ let RenderMethod = "addMemOperands", SuperClasses = [X86MemAsmOperand] in {
   def X86Mem512_RC256XOperand : AsmOperandClass { let Name = "Mem512_RC256X"; }
   def X86Mem256_RC512Operand  : AsmOperandClass { let Name = "Mem256_RC512"; }
   def X86Mem512_RC512Operand  : AsmOperandClass { let Name = "Mem512_RC512"; }
+  def X86Mem512_GR16Operand : AsmOperandClass { let Name = "Mem512_GR16"; }
+  def X86Mem512_GR32Operand : AsmOperandClass { let Name = "Mem512_GR32"; }
+  def X86Mem512_GR64Operand : AsmOperandClass { let Name = "Mem512_GR64"; }
 
   def X86SibMemOperand : AsmOperandClass { let Name = "SibMem"; }
 }
@@ -432,6 +435,11 @@ def f128mem : X86MemOperand<"printxmmwordmem", X86Mem128AsmOperand, 128>;
 def f256mem : X86MemOperand<"printymmwordmem", X86Mem256AsmOperand, 256>;
 def f512mem : X86MemOperand<"printzmmwordmem", X86Mem512AsmOperand, 512>;
 
+// 32/64 mode specific mem operands
+def i512mem_GR16 : X86MemOperand<"printzmmwordmem", X86Mem512_GR16Operand, 512>;
+def i512mem_GR32 : X86MemOperand<"printzmmwordmem", X86Mem512_GR32Operand, 512>;
+def i512mem_GR64 : X86MemOperand<"printzmmwordmem", X86Mem512_GR64Operand, 512>;
+
 // Gather mem operands
 def vx64mem  : X86VMemOperand<VR128,  "printqwordmem",  X86Mem64_RC128Operand, 64>;
 def vx128mem : X86VMemOperand<VR128,  "printxmmwordmem", X86Mem128_RC128Operand, 128>;
diff --git a/llvm/lib/Target/X86/X86InstrMisc.td b/llvm/lib/Target/X86/X86InstrMisc.td
index 092a10812d60..bf094f3ef2a3 100644
--- a/llvm/lib/Target/X86/X86InstrMisc.td
+++ b/llvm/lib/Target/X86/X86InstrMisc.td
@@ -1521,14 +1521,14 @@ def MOVDIRI64 : RI<0xF9, MRMDestMem, (outs), (ins i64mem:$dst, GR64:$src),
 // MOVDIR64B - Move 64 bytes as direct store
 //
 let SchedRW = [WriteStore] in {
-def MOVDIR64B16 : I<0xF8, MRMSrcMem, (outs), (ins GR16:$dst, i512mem:$src),
+def MOVDIR64B16 : I<0xF8, MRMSrcMem, (outs), (ins GR16:$dst, i512mem_GR16:$src),
                     "movdir64b\t{$src, $dst|$dst, $src}", []>,
                    T8PD, AdSize16, Requires<[HasMOVDIR64B, Not64BitMode]>;
-def MOVDIR64B32 : I<0xF8, MRMSrcMem, (outs), (ins GR32:$dst, i512mem:$src),
+def MOVDIR64B32 : I<0xF8, MRMSrcMem, (outs), (ins GR32:$dst, i512mem_GR32:$src),
                     "movdir64b\t{$src, $dst|$dst, $src}",
                     [(int_x86_movdir64b GR32:$dst, addr:$src)]>,
                    T8PD, AdSize32, Requires<[HasMOVDIR64B]>;
-def MOVDIR64B64 : I<0xF8, MRMSrcMem, (outs), (ins GR64:$dst, i512mem:$src),
+def MOVDIR64B64 : I<0xF8, MRMSrcMem, (outs), (ins GR64:$dst, i512mem_GR64:$src),
                     "movdir64b\t{$src, $dst|$dst, $src}",
                     [(int_x86_movdir64b GR64:$dst, addr:$src)]>,
                    T8PD, AdSize64, Requires<[HasMOVDIR64B, In64BitMode]>;
diff --git a/llvm/test/MC/X86/index-operations.s b/llvm/test/MC/X86/index-operations.s
index 0c39fe8a90da..a355b7ae0760 100644
--- a/llvm/test/MC/X86/index-operations.s
+++ b/llvm/test/MC/X86/index-operations.s
@@ -160,3 +160,29 @@ insw %dx, (%rbx)
 // ERR32: 64-bit
 // ERR16: 64-bit
 
+movdir64b	291(%si), %ecx
+// ERR32: invalid operand
+// ERR16: invalid operand
+
+movdir64b	291(%esi), %cx
+// ERR32: invalid operand
+// ERR16: invalid operand
+
+movdir64b (%rdx), %r15d
+// ERR64: invalid operand
+
+movdir64b (%edx), %r15
+// ERR64: invalid operand
+
+movdir64b (%eip), %ebx
+// 64: movdir64b (%eip), %ebx # encoding: [0x67,0x66,0x0f,0x38,0xf8,0x1d,0x00,0x00,0x00,0x00]
+
+movdir64b (%rip), %rbx
+// 64: movdir64b (%rip), %rbx # encoding: [0x66,0x0f,0x38,0xf8,0x1d,0x00,0x00,0x00,0x00]
+
+movdir64b 291(%esi, %eiz, 4), %ebx
+// 64: movdir64b 291(%esi,%eiz,4), %ebx # encoding: [0x67,0x66,0x0f,0x38,0xf8,0x9c,0xa6,0x23,0x01,0x00,0x00]
+// 32: movdir64b 291(%esi,%eiz,4), %ebx # encoding: [0x66,0x0f,0x38,0xf8,0x9c,0xa6,0x23,0x01,0x00,0x00]
+
+movdir64b 291(%rsi, %riz, 4), %rbx
+// 64: movdir64b 291(%rsi,%riz,4), %rbx # encoding: [0x66,0x0f,0x38,0xf8,0x9c,0xa6,0x23,0x01,0x00,0x00]
\ No newline at end of file
diff --git a/llvm/utils/TableGen/X86RecognizableInstr.cpp b/llvm/utils/TableGen/X86RecognizableInstr.cpp
index e5c1e53936f6..6aeb200862db 100644
--- a/llvm/utils/TableGen/X86RecognizableInstr.cpp
+++ b/llvm/utils/TableGen/X86RecognizableInstr.cpp
@@ -955,6 +955,9 @@ OperandType RecognizableInstr::typeFromString(const std::string &s,
   TYPE("i128mem",             TYPE_M)
   TYPE("i256mem",             TYPE_M)
   TYPE("i512mem",             TYPE_M)
+  TYPE("i512mem_GR16",        TYPE_M)
+  TYPE("i512mem_GR32",        TYPE_M)
+  TYPE("i512mem_GR64",        TYPE_M)
   TYPE("i64i32imm_brtarget",  TYPE_REL)
   TYPE("i16imm_brtarget",     TYPE_REL)
   TYPE("i32imm_brtarget",     TYPE_REL)
@@ -1221,6 +1224,9 @@ RecognizableInstr::memoryEncodingFromString(const std::string &s,
   ENCODING("i128mem",         ENCODING_RM)
   ENCODING("i256mem",         ENCODING_RM)
   ENCODING("i512mem",         ENCODING_RM)
+  ENCODING("i512mem_GR16",    ENCODING_RM)
+  ENCODING("i512mem_GR32",    ENCODING_RM)
+  ENCODING("i512mem_GR64",    ENCODING_RM)
   ENCODING("f80mem",          ENCODING_RM)
   ENCODING("lea64_32mem",     ENCODING_RM)
   ENCODING("lea64mem",        ENCODING_RM)