[AArch64][GlobalISel] Legalize vector boolean bitcasts to scalars by lowering via stack.

llvm/include/llvm/CodeGen/GlobalISel/LegalizerHelper.h

@@ -297,6 +297,10 @@ class LegalizerHelper {
   MachineInstrBuilder createStackTemporary(TypeSize Bytes, Align Alignment,
                                            MachinePointerInfo &PtrInfo);
 
+  /// Create a store of \p Val to a stack temporary and return a load of the
+  /// same value as type \p DestVT.
+  MachineInstrBuilder createStackStoreLoad(Register Val, LLT DstTy);
+
   /// Get a pointer to vector element \p Index located in memory for a vector of
   /// type \p VecTy starting at a base address of \p VecPtr. If \p Index is out
   /// of bounds the returned pointer is unspecified, but will be within the

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

@@ -22,6 +22,7 @@
 #include "llvm/CodeGen/GlobalISel/MIPatternMatch.h"
 #include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
 #include "llvm/CodeGen/GlobalISel/Utils.h"
+#include "llvm/CodeGen/LowLevelTypeUtils.h"
 #include "llvm/CodeGen/MachineConstantPool.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
@@ -3022,8 +3023,18 @@ LegalizerHelper::widenScalar(MachineInstr &MI, unsigned TypeIdx, LLT WideTy) {
       return UnableToLegalize;
 
     LLT Ty = MRI.getType(MI.getOperand(0).getReg());
-    if (!Ty.isScalar())
-      return UnableToLegalize;
+    if (!Ty.isScalar()) {
+      // We need to widen the vector element type.
+      Observer.changingInstr(MI);
+      widenScalarSrc(MI, WideTy, 0, TargetOpcode::G_ANYEXT);
+      // We also need to adjust the MMO to turn this into a truncating store.
+      MachineMemOperand &MMO = **MI.memoperands_begin();
+      MachineFunction &MF = MIRBuilder.getMF();
+      auto *NewMMO = MF.getMachineMemOperand(&MMO, MMO.getPointerInfo(), Ty);
+      MI.setMemRefs(MF, {NewMMO});
+      Observer.changedInstr(MI);
+      return Legalized;
+    }
 
     Observer.changingInstr(MI);
 
@@ -4653,6 +4664,20 @@ LegalizerHelper::createStackTemporary(TypeSize Bytes, Align Alignment,
   return MIRBuilder.buildFrameIndex(FramePtrTy, FrameIdx);
 }
 
+MachineInstrBuilder LegalizerHelper::createStackStoreLoad(Register Val,
+                                                          LLT DstTy) {
+  LLT SrcTy = MRI.getType(Val);
+  Align StackTypeAlign = getStackTemporaryAlignment(SrcTy);
+  MachinePointerInfo PtrInfo;
+  auto StackTemp =
+      createStackTemporary(SrcTy.getSizeInBytes(), StackTypeAlign, PtrInfo);
+
+  MIRBuilder.buildStore(Val, StackTemp, PtrInfo, StackTypeAlign);
+  return MIRBuilder.buildLoad(
+      DstTy, StackTemp, PtrInfo,
+      std::min(StackTypeAlign, getStackTemporaryAlignment(DstTy)));
+}
+
 static Register clampVectorIndex(MachineIRBuilder &B, Register IdxReg,
                                  LLT VecTy) {
   LLT IdxTy = B.getMRI()->getType(IdxReg);

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

@@ -454,6 +454,7 @@ AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
           {nxv2s64, p0, nxv2s64, 8},
       })
       .clampScalar(0, s8, s64)
+      .minScalarOrElt(0, s8)
       .lowerIf([=](const LegalityQuery &Query) {
         return Query.Types[0].isScalar() &&
                Query.Types[0] != Query.MMODescrs[0].MemoryTy;
@@ -861,6 +862,13 @@ AArch64LegalizerInfo::AArch64LegalizerInfo(const AArch64Subtarget &ST)
       .legalForCartesianProduct({s32, v2s16, v4s8})
       .legalForCartesianProduct({s64, v8s8, v4s16, v2s32})
       .legalForCartesianProduct({s128, v16s8, v8s16, v4s32, v2s64, v2p0})
+      .customIf([=](const LegalityQuery &Query) {
+        // Handle casts from i1 vectors to scalars.
+        LLT DstTy = Query.Types[0];
+        LLT SrcTy = Query.Types[1];
+        return DstTy.isScalar() && SrcTy.isVector() &&
+               SrcTy.getScalarSizeInBits() == 1;
+      })
       .lowerIf([=](const LegalityQuery &Query) {
         return Query.Types[0].isVector() != Query.Types[1].isVector();
       })
@@ -1404,11 +1412,29 @@ bool AArch64LegalizerInfo::legalizeCustom(
     return Helper.lowerAbsToCNeg(MI);
   case TargetOpcode::G_ICMP:
     return legalizeICMP(MI, MRI, MIRBuilder);
+  case TargetOpcode::G_BITCAST:
+    return legalizeBitcast(MI, Helper);
   }
 
   llvm_unreachable("expected switch to return");
 }
 
+bool AArch64LegalizerInfo::legalizeBitcast(MachineInstr &MI,
+                                           LegalizerHelper &Helper) const {
+  assert(MI.getOpcode() == TargetOpcode::G_BITCAST && "Unexpected opcode");
+  auto [DstReg, DstTy, SrcReg, SrcTy] = MI.getFirst2RegLLTs();
+  // We're trying to handle casts from i1 vectors to scalars but reloading from
+  // stack.
+  if (!DstTy.isScalar() || !SrcTy.isVector() ||
+      SrcTy.getElementType() != LLT::scalar(1))
+    return false;
+
+  auto Load = Helper.createStackStoreLoad(SrcReg, DstTy);
+  Helper.MIRBuilder.buildCopy(DstReg, Load.getReg(0));
+  MI.eraseFromParent();
+  return true;
+}
+
 bool AArch64LegalizerInfo::legalizeFunnelShift(MachineInstr &MI,
                                                MachineRegisterInfo &MRI,
                                                MachineIRBuilder &MIRBuilder,

llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.h

@@ -66,6 +66,7 @@ class AArch64LegalizerInfo : public LegalizerInfo {
                                 LegalizerHelper &Helper) const;
   bool legalizeDynStackAlloc(MachineInstr &MI, LegalizerHelper &Helper) const;
   bool legalizePrefetch(MachineInstr &MI, LegalizerHelper &Helper) const;
+  bool legalizeBitcast(MachineInstr &MI, LegalizerHelper &Helper) const;
   const AArch64Subtarget *ST;
 };
 } // End llvm namespace.

llvm/test/CodeGen/AArch64/GlobalISel/legalize-bitcast.mir

@@ -1,5 +1,5 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 4
-# RUN: llc -mtriple=aarch64 -run-pass=legalizer -global-isel-abort=1  %s -o - | FileCheck %s
+# RUN: llc -mtriple=aarch64 -run-pass=legalizer -global-isel-abort=2  %s -o - | FileCheck %s
 ---
 name:            scalar_to_oversize_vector
 tracksRegLiveness: true
@@ -48,3 +48,25 @@ body:             |
     G_BR %bb.2
 
 ...
+# This test currently is expected to fall back after reaching truncstore of <8 x s8> as <8 x s1>.
+---
+name:            boolean_vector_to_scalar
+tracksRegLiveness: true
+body:             |
+    bb.1:
+    ; CHECK-LABEL: name: boolean_vector_to_scalar
+    ; CHECK: %vec:_(<8 x s1>) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[FRAME_INDEX:%[0-9]+]]:_(p0) = G_FRAME_INDEX %stack.0
+    ; CHECK-NEXT: [[ANYEXT:%[0-9]+]]:_(<8 x s8>) = G_ANYEXT %vec(<8 x s1>)
+    ; CHECK-NEXT: G_STORE [[ANYEXT]](<8 x s8>), [[FRAME_INDEX]](p0) :: (store (<8 x s1>) into %stack.0)
+    ; CHECK-NEXT: [[LOAD:%[0-9]+]]:_(s8) = G_LOAD [[FRAME_INDEX]](p0) :: (load (s8) from %stack.0)
+    ; CHECK-NEXT: %bc:_(s8) = COPY [[LOAD]](s8)
+    ; CHECK-NEXT: %ext:_(s32) = G_ANYEXT %bc(s8)
+    ; CHECK-NEXT: $w0 = COPY %ext(s32)
+    ; CHECK-NEXT: RET_ReallyLR implicit $w0
+      %vec:_(<8 x s1>) = G_IMPLICIT_DEF
+      %bc:_(s8) = G_BITCAST %vec(<8 x s1>)
+      %ext:_(s32) = G_ANYEXT %bc(s8)
+      $w0 = COPY %ext(s32)
+      RET_ReallyLR implicit $w0
+...

llvm/test/CodeGen/AArch64/GlobalISel/legalize-store-vector-bools.mir

@@ -0,0 +1,32 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
+# RUN: llc -O0 -mtriple=aarch64 -run-pass=legalizer -global-isel-abort=2 %s -o - | FileCheck %s
+# This test currently is expected to fall back after reaching truncstore of <8 x s8> as <8 x s1>.
+---
+name:            store_8xs1
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $q0, $q1, $x0
+    ; CHECK-LABEL: name: store_8xs1
+    ; CHECK: liveins: $q0, $q1, $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<4 x s32>) = COPY $q0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(<4 x s32>) = COPY $q1
+    ; CHECK-NEXT: %ptr:_(p0) = COPY $x0
+    ; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<8 x s32>) = G_CONCAT_VECTORS [[COPY]](<4 x s32>), [[COPY1]](<4 x s32>)
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+    ; CHECK-NEXT: [[BUILD_VECTOR:%[0-9]+]]:_(<8 x s32>) = G_BUILD_VECTOR [[C]](s32), [[C]](s32), [[C]](s32), [[C]](s32), [[C]](s32), [[C]](s32), [[C]](s32), [[C]](s32)
+    ; CHECK-NEXT: [[ICMP:%[0-9]+]]:_(<8 x s1>) = G_ICMP intpred(slt), [[CONCAT_VECTORS]](<8 x s32>), [[BUILD_VECTOR]]
+    ; CHECK-NEXT: [[ANYEXT:%[0-9]+]]:_(<8 x s8>) = G_ANYEXT [[ICMP]](<8 x s1>)
+    ; CHECK-NEXT: G_STORE [[ANYEXT]](<8 x s8>), %ptr(p0) :: (store (<8 x s1>))
+    ; CHECK-NEXT: RET_ReallyLR
+    %1:_(<4 x s32>) = COPY $q0
+    %2:_(<4 x s32>) = COPY $q1
+    %ptr:_(p0) = COPY $x0
+    %0:_(<8 x s32>) = G_CONCAT_VECTORS %1(<4 x s32>), %2(<4 x s32>)
+    %4:_(s32) = G_CONSTANT i32 0
+    %3:_(<8 x s32>) = G_BUILD_VECTOR %4(s32), %4(s32), %4(s32), %4(s32), %4(s32), %4(s32), %4(s32), %4(s32)
+    %5:_(<8 x s1>) = G_ICMP intpred(slt), %0(<8 x s32>), %3
+    G_STORE %5(<8 x s1>), %ptr(p0) :: (store (<8 x s1>))
+    RET_ReallyLR
+...