[X86] Added support for 8 and 16bit LEA instructions (#122102)

JaydeepChauhan14 · KornevNikita · commit 73e6736e9517 · 2025-05-21T15:59:08.000+02:00
diff --git a/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp b/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
@@ -226,9 +226,8 @@ namespace {
     bool selectLEAAddr(SDValue N, SDValue &Base,
                        SDValue &Scale, SDValue &Index, SDValue &Disp,
                        SDValue &Segment);
-    bool selectLEA64_32Addr(SDValue N, SDValue &Base,
-                            SDValue &Scale, SDValue &Index, SDValue &Disp,
-                            SDValue &Segment);
+    bool selectLEA64_Addr(SDValue N, SDValue &Base, SDValue &Scale,
+                          SDValue &Index, SDValue &Disp, SDValue &Segment);
     bool selectTLSADDRAddr(SDValue N, SDValue &Base,
                            SDValue &Scale, SDValue &Index, SDValue &Disp,
                            SDValue &Segment);
@@ -3058,36 +3057,46 @@ bool X86DAGToDAGISel::selectMOV64Imm32(SDValue N, SDValue &Imm) {
   return !TM.isLargeGlobalValue(GV);
 }
 
-bool X86DAGToDAGISel::selectLEA64_32Addr(SDValue N, SDValue &Base,
-                                         SDValue &Scale, SDValue &Index,
-                                         SDValue &Disp, SDValue &Segment) {
+bool X86DAGToDAGISel::selectLEA64_Addr(SDValue N, SDValue &Base, SDValue &Scale,
+                                       SDValue &Index, SDValue &Disp,
+                                       SDValue &Segment) {
   // Save the debug loc before calling selectLEAAddr, in case it invalidates N.
   SDLoc DL(N);
 
   if (!selectLEAAddr(N, Base, Scale, Index, Disp, Segment))
     return false;
 
+  EVT BaseType = Base.getValueType();
+  unsigned SubReg;
+  if (BaseType == MVT::i8)
+    SubReg = X86::sub_8bit;
+  else if (BaseType == MVT::i16)
+    SubReg = X86::sub_16bit;
+  else
+    SubReg = X86::sub_32bit;
+
   auto *RN = dyn_cast<RegisterSDNode>(Base);
   if (RN && RN->getReg() == 0)
     Base = CurDAG->getRegister(0, MVT::i64);
-  else if (Base.getValueType() == MVT::i32 && !isa<FrameIndexSDNode>(Base)) {
+  else if ((BaseType == MVT::i8 || BaseType == MVT::i16 ||
+            BaseType == MVT::i32) &&
+           !isa<FrameIndexSDNode>(Base)) {
     // Base could already be %rip, particularly in the x32 ABI.
     SDValue ImplDef = SDValue(CurDAG->getMachineNode(X86::IMPLICIT_DEF, DL,
                                                      MVT::i64), 0);
-    Base = CurDAG->getTargetInsertSubreg(X86::sub_32bit, DL, MVT::i64, ImplDef,
-                                         Base);
+    Base = CurDAG->getTargetInsertSubreg(SubReg, DL, MVT::i64, ImplDef, Base);
   }
 
+  EVT IndexType = Index.getValueType();
   RN = dyn_cast<RegisterSDNode>(Index);
   if (RN && RN->getReg() == 0)
     Index = CurDAG->getRegister(0, MVT::i64);
   else {
-    assert(Index.getValueType() == MVT::i32 &&
-           "Expect to be extending 32-bit registers for use in LEA");
+    assert((IndexType == BaseType) &&
+           "Expect to be extending 8/16/32-bit registers for use in LEA");
     SDValue ImplDef = SDValue(CurDAG->getMachineNode(X86::IMPLICIT_DEF, DL,
                                                      MVT::i64), 0);
-    Index = CurDAG->getTargetInsertSubreg(X86::sub_32bit, DL, MVT::i64, ImplDef,
-                                          Index);
+    Index = CurDAG->getTargetInsertSubreg(SubReg, DL, MVT::i64, ImplDef, Index);
   }
 
   return true;
diff --git a/llvm/lib/Target/X86/X86InstrArithmetic.td b/llvm/lib/Target/X86/X86InstrArithmetic.td
@@ -25,11 +25,25 @@ let SchedRW = [WriteLEA] in {
                      [(set GR32:$dst, lea32addr:$src)]>,
                      OpSize32, Requires<[Not64BitMode]>;
 
-  def LEA64_32r : I<0x8D, MRMSrcMem,
-                    (outs GR32:$dst), (ins lea64_32mem:$src),
+  let Predicates = [HasNDD], isCodeGenOnly = 1 in {
+    def LEA64_8r : I<0x8D, MRMSrcMem, (outs GR8:$dst), (ins lea64_8mem:$src),
+                     "lea{w}\t{$src|$dst}, {$dst|$src}",
+                     [(set GR8:$dst, lea64_iaddr:$src)]>,
+                   OpSize16,
+                   Requires<[In64BitMode]>;
+
+    def LEA64_16r : I<0x8D, MRMSrcMem, (outs GR16:$dst), (ins lea64_16mem:$src),
+                      "lea{w}\t{$src|$dst}, {$dst|$src}",
+                      [(set GR16:$dst, lea64_iaddr:$src)]>,
+                    OpSize16,
+                    Requires<[In64BitMode]>;
+  }
+
+  def LEA64_32r : I<0x8D, MRMSrcMem, (outs GR32:$dst), (ins lea64_32mem:$src),
                     "lea{l}\t{$src|$dst}, {$dst|$src}",
-                    [(set GR32:$dst, lea64_32addr:$src)]>,
-                    OpSize32, Requires<[In64BitMode]>;
+                    [(set GR32:$dst, lea64_iaddr:$src)]>,
+                  OpSize32,
+                  Requires<[In64BitMode]>;
 
   let isReMaterializable = 1 in
     def LEA64r   : RI<0x8D, MRMSrcMem, (outs GR64:$dst), (ins lea64mem:$src),
diff --git a/llvm/lib/Target/X86/X86InstrFragments.td b/llvm/lib/Target/X86/X86InstrFragments.td
@@ -357,11 +357,11 @@ def addr      : ComplexPattern<iPTR, 5, "selectAddr">;
 def lea32addr : ComplexPattern<i32, 5, "selectLEAAddr",
                                [add, sub, mul, X86mul_imm, shl, or, xor, frameindex],
                                []>;
-// In 64-bit mode 32-bit LEAs can use RIP-relative addressing.
-def lea64_32addr : ComplexPattern<i32, 5, "selectLEA64_32Addr",
-                                  [add, sub, mul, X86mul_imm, shl, or, xor,
-                                   frameindex, X86WrapperRIP],
-                                  []>;
+// In 64-bit mode 8/16/32-bit LEAs can use RIP-relative addressing.
+def lea64_iaddr : ComplexPattern<iAny, 5, "selectLEA64_Addr",
+                                 [add, sub, mul, X86mul_imm, shl, or, xor,
+                                  frameindex, X86WrapperRIP],
+                                 []>;
 
 def tls32addr : ComplexPattern<i32, 5, "selectTLSADDRAddr",
                                [tglobaltlsaddr], []>;
diff --git a/llvm/lib/Target/X86/X86InstrOperands.td b/llvm/lib/Target/X86/X86InstrOperands.td
@@ -456,6 +456,18 @@ def i64u8imm : Operand<i64> {
   let OperandType = "OPERAND_IMMEDIATE";
 }
 
+def lea64_8mem : Operand<i8> {
+  let PrintMethod = "printMemReference";
+  let MIOperandInfo = (ops GR64, i8imm, GR64_NOSP, i8imm, SEGMENT_REG);
+  let ParserMatchClass = X86MemAsmOperand;
+}
+
+def lea64_16mem : Operand<i16> {
+  let PrintMethod = "printMemReference";
+  let MIOperandInfo = (ops GR64, i8imm, GR64_NOSP, i16imm, SEGMENT_REG);
+  let ParserMatchClass = X86MemAsmOperand;
+}
+
 def lea64_32mem : Operand<i32> {
   let PrintMethod = "printMemReference";
   let MIOperandInfo = (ops GR64, i8imm, GR64_NOSP, i32imm, SEGMENT_REG);
diff --git a/llvm/test/CodeGen/X86/lea-16bit.ll b/llvm/test/CodeGen/X86/lea-16bit.ll
@@ -0,0 +1,21 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -mtriple=x86_64-linux | FileCheck %s --check-prefix=NO-NDD
+; RUN: llc < %s -mtriple=x86_64-linux -mattr=+ndd | FileCheck %s --check-prefix=NDD
+
+define i16 @lea16bit(i16 %in) {
+; NO-NDD-LABEL: lea16bit:
+; NO-NDD:       # %bb.0:
+; NO-NDD-NEXT:    # kill: def $edi killed $edi def $rdi
+; NO-NDD-NEXT:    leal 1(%rdi,%rdi), %eax
+; NO-NDD-NEXT:    # kill: def $ax killed $ax killed $eax
+; NO-NDD-NEXT:    retq
+;
+; NDD-LABEL: lea16bit:
+; NDD:       # %bb.0:
+; NDD-NEXT:    # kill: def $edi killed $edi def $rdi
+; NDD-NEXT:    leaw 1(%rdi,%rdi), %ax
+; NDD-NEXT:    retq
+  %shl = shl i16 %in, 1
+  %or = or i16 %shl, 1
+  ret i16 %or
+}
diff --git a/llvm/test/CodeGen/X86/lea-8bit.ll b/llvm/test/CodeGen/X86/lea-8bit.ll
@@ -0,0 +1,22 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -mtriple=x86_64-linux | FileCheck %s --check-prefix=NO-NDD
+; RUN: llc < %s -mtriple=x86_64-linux -mattr=+ndd | FileCheck %s --check-prefix=NDD
+
+define i8 @lea8bit(i8 %in) {
+; NO-NDD-LABEL: lea8bit:
+; NO-NDD:       # %bb.0:
+; NO-NDD-NEXT:    # kill: def $edi killed $edi def $rdi
+; NO-NDD-NEXT:    leal (%rdi,%rdi), %eax
+; NO-NDD-NEXT:    incb %al
+; NO-NDD-NEXT:    # kill: def $al killed $al killed $eax
+; NO-NDD-NEXT:    retq
+;
+; NDD-LABEL: lea8bit:
+; NDD:       # %bb.0:
+; NDD-NEXT:    # kill: def $edi killed $edi def $rdi
+; NDD-NEXT:    leaw 1(%rdi,%rdi), %al
+; NDD-NEXT:    retq
+  %shl = shl i8 %in, 1
+  %or = or i8 %shl, 1
+  ret i8 %or
+}
diff --git a/llvm/utils/TableGen/X86RecognizableInstr.cpp b/llvm/utils/TableGen/X86RecognizableInstr.cpp
@@ -1096,6 +1096,8 @@ OperandType RecognizableInstr::typeFromString(const std::string &s,
   TYPE("brtarget16", TYPE_REL)
   TYPE("brtarget8", TYPE_REL)
   TYPE("f80mem", TYPE_M)
+  TYPE("lea64_8mem", TYPE_M)
+  TYPE("lea64_16mem", TYPE_M)
   TYPE("lea64_32mem", TYPE_M)
   TYPE("lea64mem", TYPE_M)
   TYPE("VR64", TYPE_MM64)
@@ -1364,6 +1366,8 @@ RecognizableInstr::memoryEncodingFromString(const std::string &s,
   ENCODING("i512mem_GR32", ENCODING_RM)
   ENCODING("i512mem_GR64", ENCODING_RM)
   ENCODING("f80mem", ENCODING_RM)
+  ENCODING("lea64_8mem", ENCODING_RM)
+  ENCODING("lea64_16mem", ENCODING_RM)
   ENCODING("lea64_32mem", ENCODING_RM)
   ENCODING("lea64mem", ENCODING_RM)
   ENCODING("anymem", ENCODING_RM)
