[SelectionDAG] Fix assertion failure on inline asm register type mismatch

Author: ahmednoursphinx

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -365,7 +365,17 @@ static SDValue getCopyFromPartsVector(SelectionDAG &DAG, const SDLoc &DL,
 
     assert(NumRegs == NumParts && "Part count doesn't match vector breakdown!");
     NumParts = NumRegs; // Silence a compiler warning.
-    assert(RegisterVT == PartVT && "Part type doesn't match vector breakdown!");
+    
+    // Check if the register type matches the part type
+    if (RegisterVT != PartVT) {
+      diagnosePossiblyInvalidConstraint(
+          *DAG.getContext(), V,
+          "register type (" + EVT(RegisterVT).getEVTString() +
+              ") doesn't match operand type (" + EVT(PartVT).getEVTString() +
+              ")");
+      return DAG.getUNDEF(ValueVT);
+    }
+    
     assert(RegisterVT.getSizeInBits() ==
            Parts[0].getSimpleValueType().getSizeInBits() &&
            "Part type sizes don't match!");

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -61825,6 +61825,8 @@ X86TargetLowering::getRegForInlineAsmConstraint(const TargetRegisterInfo *TRI,
       break;
     case 'y':   // MMX_REGS if MMX allowed.
       if (!Subtarget.hasMMX()) break;
+      // MMX registers are 64-bit only
+      if (VT.getSizeInBits() != 64) break;
       return std::make_pair(0U, &X86::VR64RegClass);
     case 'v':
     case 'x':   // SSE_REGS if SSE1 allowed or AVX_REGS if AVX allowed

llvm/test/CodeGen/X86/inline-asm-y-constraint-size-mismatch.ll

@@ -0,0 +1,20 @@
+; RUN: not llc -mtriple=x86_64-unknown-linux-gnu < %s 2>&1 | FileCheck %s
+
+; Test that using MMX register constraint 'y' (64-bit) with a 256-bit vector
+; produces a proper error message instead of an assertion failure.
+
+; CHECK: error: couldn't allocate output register for constraint 'y'
+
+define <8 x i32> @test_mmx_constraint_size_mismatch() {
+entry:
+  %out = tail call <8 x i32> asm "something $0", "=y"()
+  ret <8 x i32> %out
+}
+
+; Also test with a different vector size
+define <4 x i32> @test_mmx_constraint_128bit() {
+entry:
+  %out = tail call <4 x i32> asm "something $0", "=y"()
+  ret <4 x i32> %out
+}
+