[llvm] [LoongArch] Relax the restrictions of inlineasm operand modifier 'u' and 'w' (PR #129864)

[Lu Weining via llvm-commits]

https://github.com/SixWeining updated https://github.com/llvm/llvm-project/pull/129864

>From 32d170a4c45a1b43efc36bf20eec8eee1b4b7ea5 Mon Sep 17 00:00:00 2001
From: Weining Lu <luweining at loongson.cn>
Date: Wed, 5 Mar 2025 17:57:35 +0800
Subject: [PATCH] [LoongArch] Relax the restrictions of inlineasm operand
 modifier 'u' and 'w'

- Allow 'u' and 'w' on LASX, LSX or floating point register operands.
- Also add missing description in LangRef.

Fixes #129863.
---
 llvm/docs/LangRef.rst                         |  2 +
 .../Target/LoongArch/LoongArchAsmPrinter.cpp  | 35 ++++++++++------
 .../lasx/inline-asm-operand-modifier.ll       | 40 +++++++++++++++++++
 3 files changed, 64 insertions(+), 13 deletions(-)

diff --git a/llvm/docs/LangRef.rst b/llvm/docs/LangRef.rst
index deb87365ae8d7..78f59f64ffcd1 100644
--- a/llvm/docs/LangRef.rst
+++ b/llvm/docs/LangRef.rst
@@ -5814,6 +5814,8 @@ Hexagon:
 
 LoongArch:
 
+- ``u``: Print an LASX register.
+- ``w``: Print an LSX register.
 - ``z``: Print $zero register if operand is zero, otherwise print it normally.
 
 MSP430:
diff --git a/llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp b/llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp
index 169f9568e5362..895a8e2646692 100644
--- a/llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp
+++ b/llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp
@@ -90,20 +90,29 @@ bool LoongArchAsmPrinter::PrintAsmOperand(const MachineInstr *MI, unsigned OpNo,
         return false;
       }
       break;
-    case 'w': // Print LSX registers.
-      if (MO.getReg().id() >= LoongArch::VR0 &&
-          MO.getReg().id() <= LoongArch::VR31)
-        break;
-      // The modifier is 'w' but the operand is not an LSX register; Report an
-      // unknown operand error.
-      return true;
     case 'u': // Print LASX registers.
-      if (MO.getReg().id() >= LoongArch::XR0 &&
-          MO.getReg().id() <= LoongArch::XR31)
-        break;
-      // The modifier is 'u' but the operand is not an LASX register; Report an
-      // unknown operand error.
-      return true;
+    case 'w': // Print LSX registers.
+    {
+      // If the operand is an LASX, LSX or floating point register, print the
+      // name of LASX or LSX register with the same index in that register
+      // class.
+      unsigned RegID = MO.getReg().id(), FirstReg;
+      if (RegID >= LoongArch::XR0 && RegID <= LoongArch::XR31)
+        FirstReg = LoongArch::XR0;
+      else if (RegID >= LoongArch::VR0 && RegID <= LoongArch::VR31)
+        FirstReg = LoongArch::VR0;
+      else if (RegID >= LoongArch::F0_64 && RegID <= LoongArch::F31_64)
+        FirstReg = LoongArch::F0_64;
+      else if (RegID >= LoongArch::F0 && RegID <= LoongArch::F31)
+        FirstReg = LoongArch::F0;
+      else
+        return true;
+      OS << '$'
+         << LoongArchInstPrinter::getRegisterName(
+                RegID - FirstReg +
+                (ExtraCode[0] == 'u' ? LoongArch::XR0 : LoongArch::VR0));
+      return false;
+    }
       // TODO: handle other extra codes if any.
     }
   }
diff --git a/llvm/test/CodeGen/LoongArch/lasx/inline-asm-operand-modifier.ll b/llvm/test/CodeGen/LoongArch/lasx/inline-asm-operand-modifier.ll
index 201e34c8b5ae0..8b25a6525381b 100644
--- a/llvm/test/CodeGen/LoongArch/lasx/inline-asm-operand-modifier.ll
+++ b/llvm/test/CodeGen/LoongArch/lasx/inline-asm-operand-modifier.ll
@@ -12,3 +12,43 @@ entry:
   %0 = tail call <4 x i64> asm sideeffect "xvldi ${0:u}, 1", "=f"()
   ret void
 }
+
+define void @test_u_2xi64() nounwind {
+; CHECK-LABEL: test_u_2xi64:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    #APP
+; CHECK-NEXT:    xvldi $xr0, 1
+; CHECK-NEXT:    #NO_APP
+; CHECK-NEXT:    ret
+entry:
+  %0 = tail call <2 x i64> asm sideeffect "xvldi ${0:u}, 1", "=f"()
+  ret void
+}
+
+define void @test_w_4xi64() nounwind {
+; CHECK-LABEL: test_w_4xi64:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    #APP
+; CHECK-NEXT:    vldi $vr0, 1
+; CHECK-NEXT:    #NO_APP
+; CHECK-NEXT:    ret
+entry:
+  %0 = tail call <4 x i64> asm sideeffect "vldi ${0:w}, 1", "=f"()
+  ret void
+}
+
+define void @m128i_to_m256i(ptr %out, ptr %in) nounwind {
+; CHECK-LABEL: m128i_to_m256i:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vld $vr0, $a1, 0
+; CHECK-NEXT:    xvrepli.b $xr1, 0
+; CHECK-NEXT:    #APP
+; CHECK-NEXT:    xvpermi.q $xr1, $xr0, 32
+; CHECK-NEXT:    #NO_APP
+; CHECK-NEXT:    xvst $xr1, $a0, 0
+; CHECK-NEXT:    ret
+  %v = load <2 x i64>, ptr %in
+  %x = call <4 x i64> asm sideeffect "xvpermi.q ${0:u}, ${1:u}, 32", "=f,f,0"(<2 x i64> %v, <4 x i64> zeroinitializer)
+  store <4 x i64> %x, ptr %out
+  ret void
+}