[NVPTX] legalize v2i32 to improve codegen of v2f32 ops

Since v2f32 is legal but v2i32 is not, this causes some sequences of
operations like bitcast (build_vector) to be lowered inefficiently.

Author: Prince781

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

@@ -1018,6 +1018,7 @@ pickOpcodeForVT(MVT::SimpleValueType VT, std::optional<unsigned> Opcode_i16,
   case MVT::f32:
     return Opcode_i32;
   case MVT::v2f32:
+  case MVT::v2i32:
   case MVT::i64:
   case MVT::f64:
     return Opcode_i64;

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

@@ -228,17 +228,9 @@ getVectorLoweringShape(EVT VectorEVT, const NVPTXSubtarget &STI,
     return std::nullopt;
   case MVT::v4i64:
   case MVT::v4f64:
-  case MVT::v8i32:
-    // This is a "native" vector type iff the address space is global
-    // and the target supports 256-bit loads/stores
-    if (!CanLowerTo256Bit)
-      return std::nullopt;
-    LLVM_FALLTHROUGH;
   case MVT::v2i8:
-  case MVT::v2i32:
   case MVT::v2i64:
   case MVT::v2f64:
-  case MVT::v4i32:
     // This is a "native" vector type
     return std::pair(NumElts, EltVT);
   case MVT::v16f16:  // <8 x f16x2>
@@ -264,6 +256,7 @@ getVectorLoweringShape(EVT VectorEVT, const NVPTXSubtarget &STI,
   case MVT::v16i8:  // <4 x i8x4>
     PackRegSize = 32;
     break;
+
   case MVT::v8f32: // <4 x f32x2>
     if (!CanLowerTo256Bit)
       return std::nullopt;
@@ -274,6 +267,17 @@ getVectorLoweringShape(EVT VectorEVT, const NVPTXSubtarget &STI,
       return std::pair(NumElts, EltVT);
     PackRegSize = 64;
     break;
+
+  case MVT::v8i32:  // <4 x i32x2>
+    if (!CanLowerTo256Bit)
+      return std::nullopt;
+    LLVM_FALLTHROUGH;
+  case MVT::v2i32:  // <1 x i32x2>
+  case MVT::v4i32:  // <2 x i32x2>
+    if (!STI.hasF32x2Instructions())
+      return std::pair(NumElts, EltVT);
+    PackRegSize = 64;
+    break;
   }
 
   // If we reach here, then we can pack 2 or more elements into a single 32-bit
@@ -590,8 +594,10 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
   addRegisterClass(MVT::bf16, &NVPTX::B16RegClass);
   addRegisterClass(MVT::v2bf16, &NVPTX::B32RegClass);
 
-  if (STI.hasF32x2Instructions())
+  if (STI.hasF32x2Instructions()) {
     addRegisterClass(MVT::v2f32, &NVPTX::B64RegClass);
+    addRegisterClass(MVT::v2i32, &NVPTX::B64RegClass);
+  }
 
   // Conversion to/from FP16/FP16x2 is always legal.
   setOperationAction(ISD::BUILD_VECTOR, MVT::v2f16, Custom);
@@ -628,12 +634,13 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
   setOperationAction(ISD::INSERT_VECTOR_ELT, MVT::v4i8, Custom);
   setOperationAction(ISD::VECTOR_SHUFFLE, MVT::v4i8, Custom);
 
-  // No support for these operations with v2f32.
-  setOperationAction(ISD::INSERT_VECTOR_ELT, MVT::v2f32, Expand);
-  setOperationAction(ISD::VECTOR_SHUFFLE, MVT::v2f32, Expand);
+  // No support for these operations with v2f32/v2i32
+  setOperationAction(ISD::INSERT_VECTOR_ELT, {MVT::v2f32, MVT::v2i32}, Expand);
+  setOperationAction(ISD::VECTOR_SHUFFLE, {MVT::v2f32, MVT::v2i32}, Expand);
   // Need custom lowering in case the index is dynamic.
   if (STI.hasF32x2Instructions())
-    setOperationAction(ISD::EXTRACT_VECTOR_ELT, MVT::v2f32, Custom);
+    setOperationAction(ISD::EXTRACT_VECTOR_ELT, {MVT::v2f32, MVT::v2i32},
+                       Custom);
 
   // Custom conversions to/from v2i8.
   setOperationAction(ISD::BITCAST, MVT::v2i8, Custom);
@@ -661,14 +668,13 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
   // Operations not directly supported by NVPTX.
   for (MVT VT : {MVT::bf16, MVT::f16, MVT::v2bf16, MVT::v2f16, MVT::f32,
                  MVT::v2f32, MVT::f64, MVT::i1, MVT::i8, MVT::i16, MVT::v2i16,
-                 MVT::v4i8, MVT::i32, MVT::i64}) {
+                 MVT::v4i8, MVT::i32, MVT::v2i32, MVT::i64}) {
     setOperationAction(ISD::SELECT_CC, VT, Expand);
     setOperationAction(ISD::BR_CC, VT, Expand);
   }
 
-  // Not directly supported. TLI would attempt to expand operations like
-  // FMINIMUM(v2f32) using invalid SETCC and VSELECT nodes.
-  setOperationAction(ISD::VSELECT, MVT::v2f32, Expand);
+  // We don't want ops like FMINIMUM or UMAX to be lowered to SETCC+VSELECT.
+  setOperationAction(ISD::VSELECT, {MVT::v2f32, MVT::v2i32}, Expand);
 
   // Some SIGN_EXTEND_INREG can be done using cvt instruction.
   // For others we will expand to a SHL/SRA pair.
@@ -815,7 +821,14 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
   setOperationAction({ISD::SDIV, ISD::UDIV, ISD::SRA, ISD::SRL, ISD::MULHS,
                       ISD::MULHU, ISD::FP_TO_SINT, ISD::FP_TO_UINT,
                       ISD::SINT_TO_FP, ISD::UINT_TO_FP, ISD::SETCC},
-                     MVT::v2i16, Expand);
+                     {MVT::v2i16, MVT::v2i32}, Expand);
+
+  // v2i32 is not supported for any arithmetic operations
+  setOperationAction({ISD::ABS, ISD::SMIN, ISD::SMAX, ISD::UMIN, ISD::UMAX,
+                      ISD::CTPOP, ISD::CTLZ, ISD::ADD, ISD::SUB, ISD::MUL,
+                      ISD::SHL, ISD::SRA, ISD::SRL, ISD::OR, ISD::AND, ISD::XOR,
+                      ISD::SREM, ISD::UREM},
+                     MVT::v2i32, Expand);
 
   setOperationAction(ISD::ADDC, MVT::i32, Legal);
   setOperationAction(ISD::ADDE, MVT::i32, Legal);
@@ -829,7 +842,7 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
   }
 
   setOperationAction(ISD::CTTZ, MVT::i16, Expand);
-  setOperationAction(ISD::CTTZ, MVT::v2i16, Expand);
+  setOperationAction(ISD::CTTZ, {MVT::v2i16, MVT::v2i32}, Expand);
   setOperationAction(ISD::CTTZ, MVT::i32, Expand);
   setOperationAction(ISD::CTTZ, MVT::i64, Expand);
 
@@ -5673,7 +5686,7 @@ static SDValue PerformEXTRACTCombine(SDNode *N,
       IsPTXVectorType(VectorVT.getSimpleVT()))
     return SDValue(); // Native vector loads already combine nicely w/
                       // extract_vector_elt.
-  // Don't mess with singletons or packed types (v2f32, v2*16, v4i8 and v8i8),
+  // Don't mess with singletons or packed types (v2*32, v2*16, v4i8 and v8i8),
   // we already handle them OK.
   if (VectorVT.getVectorNumElements() == 1 ||
       NVPTX::isPackedVectorTy(VectorVT) || VectorVT == MVT::v8i8)

llvm/lib/Target/NVPTX/NVPTXInstrInfo.td

@@ -754,8 +754,10 @@ def : Pat<(vt (select i1:$p, vt:$a, vt:$b)),
           (SELP_b32rr $a, $b, $p)>;
 }
 
-def : Pat<(v2f32 (select i1:$p, v2f32:$a, v2f32:$b)),
+foreach vt = [v2f32, v2i32] in {
+def : Pat<(vt (select i1:$p, vt:$a, vt:$b)),
           (SELP_b64rr $a, $b, $p)>;
+}
 
 //-----------------------------------
 // Test Instructions
@@ -2092,8 +2094,8 @@ foreach vt = [v2f16, v2bf16, v2i16] in {
             (V2I16toI32 $a, $b)>;
 }
 
-// Same thing for the 64-bit type v2f32.
-foreach vt = [v2f32] in {
+// Handle extracting one element from the pair (64-bit types)
+foreach vt = [v2f32, v2i32] in {
   def : Pat<(extractelt vt:$src, 0), (I64toI32L_Sink $src)>, Requires<[hasPTX<71>]>;
   def : Pat<(extractelt vt:$src, 1), (I64toI32H_Sink $src)>, Requires<[hasPTX<71>]>;
 

llvm/lib/Target/NVPTX/NVPTXRegisterInfo.td

@@ -54,7 +54,8 @@ def B16 : NVPTXRegClass<[i16, f16, bf16], 16, (add (sequence "RS%u", 0, 4))>;
 def B32 : NVPTXRegClass<[i32, v2f16, v2bf16, v2i16, v4i8, f32], 32,
                               (add (sequence "R%u", 0, 4),
                               VRFrame32, VRFrameLocal32)>;
-def B64 : NVPTXRegClass<[i64, v2f32, f64], 64, (add (sequence "RL%u", 0, 4),
+def B64 : NVPTXRegClass<[i64, v2i32, v2f32, f64], 64,
+                        (add (sequence "RL%u", 0, 4),
                          VRFrame64, VRFrameLocal64)>;
 // 128-bit regs are not defined as general regs in NVPTX. They are used for inlineASM only.
 def B128 : NVPTXRegClass<[i128], 128, (add (sequence "RQ%u", 0, 4))>;

llvm/lib/Target/NVPTX/NVPTXUtilities.h

@@ -99,8 +99,8 @@ namespace NVPTX {
 // register. NOTE: This must be kept in sync with the register classes
 // defined in NVPTXRegisterInfo.td.
 inline auto packed_types() {
-  static const auto PackedTypes = {MVT::v4i8, MVT::v2f16, MVT::v2bf16,
-                                   MVT::v2i16, MVT::v2f32};
+  static const auto PackedTypes = {MVT::v4i8,  MVT::v2f16, MVT::v2bf16,
+                                   MVT::v2i16, MVT::v2f32, MVT::v2i32};
   return PackedTypes;
 }
 

llvm/test/CodeGen/NVPTX/f32x2-convert-i32x2.ll

@@ -0,0 +1,119 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -mcpu=sm_90a -O0 -disable-post-ra -frame-pointer=all          \
+; RUN: -verify-machineinstrs | FileCheck --check-prefixes=CHECK,CHECK-SM90A %s
+; RUN: %if ptxas-12.7 %{                                                      \
+; RUN:  llc < %s -mcpu=sm_90a -O0 -disable-post-ra -frame-pointer=all         \
+; RUN:  -verify-machineinstrs | %ptxas-verify -arch=sm_90a                    \
+; RUN: %}
+; RUN: llc < %s -mcpu=sm_100 -O0 -disable-post-ra -frame-pointer=all          \
+; RUN: -verify-machineinstrs | FileCheck --check-prefixes=CHECK,CHECK-SM100 %s
+; RUN: %if ptxas-12.7 %{                                                      \
+; RUN:  llc < %s -mcpu=sm_100 -O0 -disable-post-ra -frame-pointer=all         \
+; RUN:  -verify-machineinstrs | %ptxas-verify -arch=sm_100                    \
+; RUN: %}
+
+; Test that v2i32 -> v2f32 conversions don't emit bitwise operations on i64.
+
+target triple = "nvptx64-nvidia-cuda"
+
+declare <2 x i32> @return_i32x2(i32 %0)
+
+; Test with v2i32.
+define ptx_kernel void @store_i32x2(i32 %0, ptr %p) {
+; CHECK-SM90A-LABEL: store_i32x2(
+; CHECK-SM90A:       {
+; CHECK-SM90A-NEXT:    .reg .b32 %r<6>;
+; CHECK-SM90A-NEXT:    .reg .b64 %rd<2>;
+; CHECK-SM90A-EMPTY:
+; CHECK-SM90A-NEXT:  // %bb.0:
+; CHECK-SM90A-NEXT:    ld.param.b64 %rd1, [store_i32x2_param_1];
+; CHECK-SM90A-NEXT:    ld.param.b32 %r1, [store_i32x2_param_0];
+; CHECK-SM90A-NEXT:    { // callseq 0, 0
+; CHECK-SM90A-NEXT:    .param .b32 param0;
+; CHECK-SM90A-NEXT:    .param .align 8 .b8 retval0[8];
+; CHECK-SM90A-NEXT:    st.param.b32 [param0], %r1;
+; CHECK-SM90A-NEXT:    call.uni (retval0), return_i32x2, (param0);
+; CHECK-SM90A-NEXT:    ld.param.v2.b32 {%r2, %r3}, [retval0];
+; CHECK-SM90A-NEXT:    } // callseq 0
+; CHECK-SM90A-NEXT:    add.rn.f32 %r4, %r3, %r3;
+; CHECK-SM90A-NEXT:    add.rn.f32 %r5, %r2, %r2;
+; CHECK-SM90A-NEXT:    st.v2.b32 [%rd1], {%r5, %r4};
+; CHECK-SM90A-NEXT:    ret;
+;
+; CHECK-SM100-LABEL: store_i32x2(
+; CHECK-SM100:       {
+; CHECK-SM100-NEXT:    .reg .b32 %r<2>;
+; CHECK-SM100-NEXT:    .reg .b64 %rd<4>;
+; CHECK-SM100-EMPTY:
+; CHECK-SM100-NEXT:  // %bb.0:
+; CHECK-SM100-NEXT:    ld.param.b64 %rd1, [store_i32x2_param_1];
+; CHECK-SM100-NEXT:    ld.param.b32 %r1, [store_i32x2_param_0];
+; CHECK-SM100-NEXT:    { // callseq 0, 0
+; CHECK-SM100-NEXT:    .param .b32 param0;
+; CHECK-SM100-NEXT:    .param .align 8 .b8 retval0[8];
+; CHECK-SM100-NEXT:    st.param.b32 [param0], %r1;
+; CHECK-SM100-NEXT:    call.uni (retval0), return_i32x2, (param0);
+; CHECK-SM100-NEXT:    ld.param.b64 %rd2, [retval0];
+; CHECK-SM100-NEXT:    } // callseq 0
+; CHECK-SM100-NEXT:    add.rn.f32x2 %rd3, %rd2, %rd2;
+; CHECK-SM100-NEXT:    st.b64 [%rd1], %rd3;
+; CHECK-SM100-NEXT:    ret;
+  %v = call <2 x i32> @return_i32x2(i32 %0)
+  %v.f32x2 = bitcast <2 x i32> %v to <2 x float>
+  %res = fadd <2 x float> %v.f32x2, %v.f32x2
+  store <2 x float> %res, ptr %p, align 8
+  ret void
+}
+
+; Test with inline ASM returning { <1 x float>, <1 x float> }, which decays to
+; v2i32.
+define ptx_kernel void @inlineasm(ptr %p) {
+; CHECK-SM90A-LABEL: inlineasm(
+; CHECK-SM90A:       {
+; CHECK-SM90A-NEXT:    .reg .b32 %r<7>;
+; CHECK-SM90A-NEXT:    .reg .b64 %rd<2>;
+; CHECK-SM90A-EMPTY:
+; CHECK-SM90A-NEXT:  // %bb.0:
+; CHECK-SM90A-NEXT:    ld.param.b64 %rd1, [inlineasm_param_0];
+; CHECK-SM90A-NEXT:    mov.b32 %r3, 0;
+; CHECK-SM90A-NEXT:    mov.b32 %r4, %r3;
+; CHECK-SM90A-NEXT:    mov.b32 %r2, %r4;
+; CHECK-SM90A-NEXT:    mov.b32 %r1, %r3;
+; CHECK-SM90A-NEXT:    // begin inline asm
+; CHECK-SM90A-NEXT:    // nop
+; CHECK-SM90A-NEXT:    // end inline asm
+; CHECK-SM90A-NEXT:    mul.rn.f32 %r5, %r2, 0f00000000;
+; CHECK-SM90A-NEXT:    mul.rn.f32 %r6, %r1, 0f00000000;
+; CHECK-SM90A-NEXT:    st.v2.b32 [%rd1], {%r6, %r5};
+; CHECK-SM90A-NEXT:    ret;
+;
+; CHECK-SM100-LABEL: inlineasm(
+; CHECK-SM100:       {
+; CHECK-SM100-NEXT:    .reg .b32 %r<6>;
+; CHECK-SM100-NEXT:    .reg .b64 %rd<5>;
+; CHECK-SM100-EMPTY:
+; CHECK-SM100-NEXT:  // %bb.0:
+; CHECK-SM100-NEXT:    ld.param.b64 %rd1, [inlineasm_param_0];
+; CHECK-SM100-NEXT:    mov.b32 %r3, 0;
+; CHECK-SM100-NEXT:    mov.b32 %r4, %r3;
+; CHECK-SM100-NEXT:    mov.b32 %r2, %r4;
+; CHECK-SM100-NEXT:    mov.b32 %r1, %r3;
+; CHECK-SM100-NEXT:    // begin inline asm
+; CHECK-SM100-NEXT:    // nop
+; CHECK-SM100-NEXT:    // end inline asm
+; CHECK-SM100-NEXT:    mov.b64 %rd2, {%r1, %r2};
+; CHECK-SM100-NEXT:    mov.b32 %r5, 0f00000000;
+; CHECK-SM100-NEXT:    mov.b64 %rd3, {%r5, %r5};
+; CHECK-SM100-NEXT:    mul.rn.f32x2 %rd4, %rd2, %rd3;
+; CHECK-SM100-NEXT:    st.b64 [%rd1], %rd4;
+; CHECK-SM100-NEXT:    ret;
+  %r = call { <1 x float>, <1 x float> } asm sideeffect "// nop", "=f,=f,0,1"(<1 x float> zeroinitializer, <1 x float> zeroinitializer)
+  %i0 = extractvalue { <1 x float>, <1 x float> } %r, 0
+  %i1 = extractvalue { <1 x float>, <1 x float> } %r, 1
+  %i4 = shufflevector <1 x float> %i0, <1 x float> %i1, <2 x i32> <i32 0, i32 1>
+  %mul = fmul < 2 x float> %i4, zeroinitializer
+  store <2 x float> %mul, ptr %p, align 8
+  ret void
+}
+;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
+; CHECK: {{.*}}