[llvm] [RISCV][GISEL] Fix legalization for G_MERGE/UNMERGE_VALUES (PR #75619)

[Michael Maitland via llvm-commits]

https://github.com/michaelmaitland updated https://github.com/llvm/llvm-project/pull/75619

>From bed69bd070c1b4d5aeade09274786f35f62e2536 Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Fri, 15 Dec 2023 07:42:39 -0800
Subject: [PATCH 1/3] [RISCV][GISEL] Fix legalization for
 G_MERGE/UNMERGE_VALUES

There are two bugs fixed by this patch:

  1. G_MERGE_VALUES is always true, causing IdxZeroTy and IdxOneTy to
  always be s64 and s32 respectively.
  2. The legalFor check was for the type in index 0, instead of for
  index 0 and index 1, so we were allowing s64 or s32 in type 0 to be legal for
  both G_MERGE_VALUES and G_UNMERGE_VALUES.
---
 llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp b/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp
index 8f03a7ac41d37b..b46269fb9565fb 100644
--- a/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp
+++ b/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp
@@ -85,14 +85,16 @@ RISCVLegalizerInfo::RISCVLegalizerInfo(const RISCVSubtarget &ST)
 
   // Merge/Unmerge
   for (unsigned Op : {G_MERGE_VALUES, G_UNMERGE_VALUES}) {
-    unsigned BigTyIdx = Op == G_MERGE_VALUES ? 0 : 1;
-    unsigned LitTyIdx = Op == G_MERGE_VALUES ? 1 : 0;
     auto &MergeUnmergeActions = getActionDefinitionsBuilder(Op);
     if (XLen == 32 && ST.hasStdExtD()) {
-      LLT IdxZeroTy = G_MERGE_VALUES ? s64 : s32;
-      LLT IdxOneTy = G_MERGE_VALUES ? s32 : s64;
-      MergeUnmergeActions.legalFor({IdxZeroTy, IdxOneTy});
+      // Only need to legalize on IdxZeroTy since there is no way to construct
+      // a valid MERGE/UNMERGE where the types don't combine/decompose as
+      // expected.
+      LLT IdxZeroTy = Op == G_MERGE_VALUES ? s64 : s32;
+      MergeUnmergeActions.legalFor({IdxZeroTy});
     }
+    unsigned BigTyIdx = Op == G_MERGE_VALUES ? 0 : 1;
+    unsigned LitTyIdx = Op == G_MERGE_VALUES ? 1 : 0;
     MergeUnmergeActions.widenScalarToNextPow2(LitTyIdx, XLen)
         .widenScalarToNextPow2(BigTyIdx, XLen)
         .clampScalar(LitTyIdx, sXLen, sXLen)

>From afb51a6455bbdbfb8bd0275cc4e3aaa7fa669d0b Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Fri, 15 Dec 2023 09:26:14 -0800
Subject: [PATCH 2/3] !fixup check both indeces

---
 llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp b/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp
index b46269fb9565fb..afade4692519bf 100644
--- a/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp
+++ b/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp
@@ -87,11 +87,9 @@ RISCVLegalizerInfo::RISCVLegalizerInfo(const RISCVSubtarget &ST)
   for (unsigned Op : {G_MERGE_VALUES, G_UNMERGE_VALUES}) {
     auto &MergeUnmergeActions = getActionDefinitionsBuilder(Op);
     if (XLen == 32 && ST.hasStdExtD()) {
-      // Only need to legalize on IdxZeroTy since there is no way to construct
-      // a valid MERGE/UNMERGE where the types don't combine/decompose as
-      // expected.
       LLT IdxZeroTy = Op == G_MERGE_VALUES ? s64 : s32;
-      MergeUnmergeActions.legalFor({IdxZeroTy});
+      LLT IdxOneTy = Op == G_MERGE_VALUES ? s32 : s64;
+      MergeUnmergeActions.legalFor({{IdxZeroTy}, {IdxOneTy}});
     }
     unsigned BigTyIdx = Op == G_MERGE_VALUES ? 0 : 1;
     unsigned LitTyIdx = Op == G_MERGE_VALUES ? 1 : 0;

>From 70dc7d7824e745d59fb3dc81e5d89456cb6c933f Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Fri, 15 Dec 2023 10:04:30 -0800
Subject: [PATCH 3/3] !fixup use legalIf

---
 llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp b/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp
index afade4692519bf..28ec999157c697 100644
--- a/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp
+++ b/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp
@@ -86,13 +86,12 @@ RISCVLegalizerInfo::RISCVLegalizerInfo(const RISCVSubtarget &ST)
   // Merge/Unmerge
   for (unsigned Op : {G_MERGE_VALUES, G_UNMERGE_VALUES}) {
     auto &MergeUnmergeActions = getActionDefinitionsBuilder(Op);
-    if (XLen == 32 && ST.hasStdExtD()) {
-      LLT IdxZeroTy = Op == G_MERGE_VALUES ? s64 : s32;
-      LLT IdxOneTy = Op == G_MERGE_VALUES ? s32 : s64;
-      MergeUnmergeActions.legalFor({{IdxZeroTy}, {IdxOneTy}});
-    }
     unsigned BigTyIdx = Op == G_MERGE_VALUES ? 0 : 1;
     unsigned LitTyIdx = Op == G_MERGE_VALUES ? 1 : 0;
+    if (XLen == 32 && ST.hasStdExtD()) {
+      MergeUnmergeActions.legalIf(
+          all(typeIs(BigTyIdx, s64), typeIs(LitTyIdx, s32)));
+    }
     MergeUnmergeActions.widenScalarToNextPow2(LitTyIdx, XLen)
         .widenScalarToNextPow2(BigTyIdx, XLen)
         .clampScalar(LitTyIdx, sXLen, sXLen)