[AArch64][GlobalISel] Add legalization for some vector G_LSR and G_ASHR

[aemerson]

Summary

This is needed for some future support for vector ICMP.

Also make G_TRUNC always legal as we need it as a legalization artifact, and the current legacy rules don't cover some vector types.

[paquette]

LGTM

llvm/lib/Target/AArch64/AArch64LegalizerInfo.cpp

277 ↗

(On Diff #194230)

Not for this patch, but in general, I wonder if we should just have a .alwaysLegal() or something.

[aemerson]

llvm/lib/Target/AArch64/AArch64LegalizerInfo.cpp

277 ↗

(On Diff #194230)

Yes, that's a good idea.

llvm/trunk/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h

public:
   /// respective lists.
   LegalizeRuleSet &legalForCartesianProduct(std::initializer_list<LLT> Types0,
                                             std::initializer_list<LLT> Types1,
                                             std::initializer_list<LLT> Types2) {
     return actionForCartesianProduct(LegalizeAction::Legal, Types0, Types1,
                                      Types2);
   }
 
+  LegalizeRuleSet &alwaysLegal() {
+    using namespace LegalizeMutations;
+    markAllTypeIdxsAsCovered();
+    return actionIf(LegalizeAction::Legal, always);
+  }
+
   /// The instruction is lowered.
   LegalizeRuleSet &lower() {
     using namespace LegalizeMutations;
     // We have no choice but conservatively assume that predicate-less lowering
     // properly handles all type indices by design:
     markAllTypeIdxsAsCovered();
     return actionIf(LegalizeAction::Lower, always);
   }

llvm/trunk/lib/Target/AArch64/AArch64LegalizerInfo.cpp

AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST) {
   getActionDefinitionsBuilder(G_PTR_MASK).legalFor({p0});
 
   getActionDefinitionsBuilder({G_SDIV, G_UDIV})
       .legalFor({s32, s64})
       .clampScalar(0, s32, s64)
       .widenScalarToNextPow2(0);
 
   getActionDefinitionsBuilder({G_LSHR, G_ASHR})
-    .legalFor({{s32, s32}, {s64, s64}})
+      .legalFor({{s32, s32}, {s64, s64}, {v2s32, v2s32}, {v4s32, v4s32}})
       .clampScalar(1, s32, s64)
       .clampScalar(0, s32, s64)
       .minScalarSameAs(1, 0);
 
   getActionDefinitionsBuilder({G_SREM, G_UREM})
       .lowerFor({s1, s8, s16, s32, s64});
 
   getActionDefinitionsBuilder({G_SMULO, G_UMULO})
       .lowerFor({{s64, s1}});
 
   getActionDefinitionsBuilder({G_SMULH, G_UMULH}).legalFor({s32, s64});
getActionDefinitionsBuilder(G_FCMP)
       .clampScalar(0, s32, s32)
       .clampScalar(1, s32, s64)
       .widenScalarToNextPow2(1);
 
   // Extensions
   getActionDefinitionsBuilder({G_ZEXT, G_SEXT, G_ANYEXT})
       .legalForCartesianProduct({s8, s16, s32, s64}, {s1, s8, s16, s32});
 
+  getActionDefinitionsBuilder(G_TRUNC).alwaysLegal();
+
   // FP conversions
   getActionDefinitionsBuilder(G_FPTRUNC).legalFor(
       {{s16, s32}, {s16, s64}, {s32, s64}, {v4s16, v4s32}, {v2s32, v2s64}});
   getActionDefinitionsBuilder(G_FPEXT).legalFor(
       {{s32, s16}, {s64, s16}, {s64, s32}, {v4s32, v4s16}, {v2s64, v2s32}});
 
   // Conversions
   getActionDefinitionsBuilder({G_FPTOSI, G_FPTOUI})

llvm/trunk/test/CodeGen/AArch64/GlobalISel/legalize-shift.mir

bb.0:
     ; CHECK: [[MV:%[0-9]+]]:_(s128) = G_MERGE_VALUES [[SELECT1]](s64), [[SELECT2]](s64)
     ; CHECK: $q0 = COPY [[MV]](s128)
     %0:_(s128) = COPY $q0
     %1:_(s128) = COPY $q1
     %2:_(s128) = G_ASHR %0, %1
     $q0 = COPY %2
 
 ...
+
+---
+name: test_ashr_v2i32
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: test_ashr_v2i32
+    ; CHECK: [[COPY:%[0-9]+]]:_(<2 x s32>) = COPY $d0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(<2 x s32>) = COPY $d1
+    ; CHECK: [[ASHR:%[0-9]+]]:_(<2 x s32>) = G_ASHR [[COPY]], [[COPY1]](<2 x s32>)
+    ; CHECK: $d0 = COPY [[ASHR]](<2 x s32>)
+    %0:_(<2 x s32>) = COPY $d0
+    %1:_(<2 x s32>) = COPY $d1
+    %2:_(<2 x s32>) = G_ASHR %0, %1
+    $d0 = COPY %2
+
+...
+---
+name: test_ashr_v4i32
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: test_ashr_v4i32
+    ; CHECK: [[COPY:%[0-9]+]]:_(<4 x s32>) = COPY $q0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(<4 x s32>) = COPY $q1
+    ; CHECK: [[ASHR:%[0-9]+]]:_(<4 x s32>) = G_ASHR [[COPY]], [[COPY1]](<4 x s32>)
+    ; CHECK: $q0 = COPY [[ASHR]](<4 x s32>)
+    %0:_(<4 x s32>) = COPY $q0
+    %1:_(<4 x s32>) = COPY $q1
+    %2:_(<4 x s32>) = G_ASHR %0, %1
+    $q0 = COPY %2
+
+...

llvm/trunk/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir

 #
 # DEBUG-NEXT: G_INTRINSIC_W_SIDE_EFFECTS (opcode {{[0-9]+}}): 0 type indices
 # DEBUG:      .. type index coverage check SKIPPED: no rules defined
 #
 # DEBUG-NEXT: G_ANYEXT (opcode {{[0-9]+}}): 2 type indices
 # DEBUG:      .. the first uncovered type index: 2, OK
 #
 # DEBUG-NEXT: G_TRUNC (opcode {{[0-9]+}}): 2 type indices
-# DEBUG:      .. type index coverage check SKIPPED: no rules defined
+# DEBUG:      .. type index coverage check SKIPPED: user-defined predicate detected
 #
 # DEBUG-NEXT: G_CONSTANT (opcode {{[0-9]+}}): 1 type index
 # DEBUG:      .. the first uncovered type index: 1, OK
 #
 # DEBUG-NEXT: G_FCONSTANT (opcode {{[0-9]+}}): 1 type index
 # DEBUG:      .. the first uncovered type index: 1, OK
 #
 # DEBUG-NEXT: G_VASTART (opcode {{[0-9]+}}): 1 type index