[llvm] [NVPTX] Set boolean contents to zero-or-one (PR #108969)

[Fraser Cormack via llvm-commits]

https://github.com/frasercrmck updated https://github.com/llvm/llvm-project/pull/108969

>From 0970633fc6d60274fed98c637c1637c58a0a4577 Mon Sep 17 00:00:00 2001
From: Fraser Cormack <fraser at codeplay.com>
Date: Tue, 17 Sep 2024 13:27:30 +0100
Subject: [PATCH] [NVPTX] Set v2i16 SETCC to Expand

Note that this refers to the return type of SETCC. This operation is not
legal in PTX but was assumed as such because v2i16 is declared a legal
type. We were already expanding v4i8 SETCC.

The DAGCombiner would in certain circumstances try to fold an extension
of an illegal v2i1 SETCC (because v2i1 is illegal) into a "legal" v2i16
SETCC, which we wouldn't have patterns for.
---
 llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp |  2 +-
 llvm/test/CodeGen/NVPTX/sext-setcc.ll       | 72 +++++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/NVPTX/sext-setcc.ll

diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
index c5a40e4308860c..be7ae1cfab3379 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -725,7 +725,7 @@ NVPTXTargetLowering::NVPTXTargetLowering(const NVPTXTargetMachine &TM,
   // Other arithmetic and logic ops are unsupported.
   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::SINT_TO_FP, ISD::UINT_TO_FP, ISD::SETCC},
                      MVT::v2i16, Expand);
 
   setOperationAction(ISD::ADDC, MVT::i32, Legal);
diff --git a/llvm/test/CodeGen/NVPTX/sext-setcc.ll b/llvm/test/CodeGen/NVPTX/sext-setcc.ll
new file mode 100644
index 00000000000000..f471d47077cf0d
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/sext-setcc.ll
@@ -0,0 +1,72 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -mtriple=nvptx64 -mcpu=sm_80 | FileCheck %s
+; RUN: %if ptxas-11.0 %{ llc < %s -march=nvptx64 -mcpu=sm_80 -mattr=+ptx70 | %ptxas-verify -arch=sm_80 %}
+
+define <2 x i16> @sext_setcc_v2i1_to_v2i16(ptr %p) {
+; CHECK-LABEL: sext_setcc_v2i1_to_v2i16(
+; CHECK:       {
+; CHECK-NEXT:    .reg .pred %p<3>;
+; CHECK-NEXT:    .reg .b16 %rs<5>;
+; CHECK-NEXT:    .reg .b32 %r<3>;
+; CHECK-NEXT:    .reg .b64 %rd<2>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0: // %entry
+; CHECK-NEXT:    ld.param.u64 %rd1, [sext_setcc_v2i1_to_v2i16_param_0];
+; CHECK-NEXT:    ld.u32 %r1, [%rd1];
+; CHECK-NEXT:    mov.b32 {%rs1, %rs2}, %r1;
+; CHECK-NEXT:    setp.eq.s16 %p1, %rs1, 0;
+; CHECK-NEXT:    setp.eq.s16 %p2, %rs2, 0;
+; CHECK-NEXT:    selp.s16 %rs3, -1, 0, %p2;
+; CHECK-NEXT:    selp.s16 %rs4, -1, 0, %p1;
+; CHECK-NEXT:    mov.b32 %r2, {%rs4, %rs3};
+; CHECK-NEXT:    st.param.b32 [func_retval0+0], %r2;
+; CHECK-NEXT:    ret;
+entry:
+  %v = load <2 x i16>, ptr %p, align 4
+  %cmp = icmp eq <2 x i16> %v, zeroinitializer
+  %sext = sext <2 x i1> %cmp to <2 x i16>
+  ret <2 x i16> %sext
+}
+
+define <4 x i8> @sext_setcc_v4i1_to_v4i8(ptr %p) {
+; CHECK-LABEL: sext_setcc_v4i1_to_v4i8(
+; CHECK:       {
+; CHECK-NEXT:    .reg .pred %p<5>;
+; CHECK-NEXT:    .reg .b16 %rs<9>;
+; CHECK-NEXT:    .reg .b32 %r<14>;
+; CHECK-NEXT:    .reg .b64 %rd<2>;
+; CHECK-EMPTY:
+; CHECK-NEXT:  // %bb.0: // %entry
+; CHECK-NEXT:    ld.param.u64 %rd1, [sext_setcc_v4i1_to_v4i8_param_0];
+; CHECK-NEXT:    ld.u32 %r1, [%rd1];
+; CHECK-NEXT:    bfe.u32 %r2, %r1, 24, 8;
+; CHECK-NEXT:    cvt.u16.u32 %rs1, %r2;
+; CHECK-NEXT:    and.b16 %rs2, %rs1, 255;
+; CHECK-NEXT:    setp.eq.s16 %p1, %rs2, 0;
+; CHECK-NEXT:    bfe.u32 %r3, %r1, 16, 8;
+; CHECK-NEXT:    cvt.u16.u32 %rs3, %r3;
+; CHECK-NEXT:    and.b16 %rs4, %rs3, 255;
+; CHECK-NEXT:    setp.eq.s16 %p2, %rs4, 0;
+; CHECK-NEXT:    bfe.u32 %r4, %r1, 8, 8;
+; CHECK-NEXT:    cvt.u16.u32 %rs5, %r4;
+; CHECK-NEXT:    and.b16 %rs6, %rs5, 255;
+; CHECK-NEXT:    setp.eq.s16 %p3, %rs6, 0;
+; CHECK-NEXT:    bfe.u32 %r5, %r1, 0, 8;
+; CHECK-NEXT:    cvt.u16.u32 %rs7, %r5;
+; CHECK-NEXT:    and.b16 %rs8, %rs7, 255;
+; CHECK-NEXT:    setp.eq.s16 %p4, %rs8, 0;
+; CHECK-NEXT:    selp.s32 %r6, -1, 0, %p4;
+; CHECK-NEXT:    selp.s32 %r7, -1, 0, %p3;
+; CHECK-NEXT:    bfi.b32 %r8, %r7, %r6, 8, 8;
+; CHECK-NEXT:    selp.s32 %r9, -1, 0, %p2;
+; CHECK-NEXT:    bfi.b32 %r10, %r9, %r8, 16, 8;
+; CHECK-NEXT:    selp.s32 %r11, -1, 0, %p1;
+; CHECK-NEXT:    bfi.b32 %r12, %r11, %r10, 24, 8;
+; CHECK-NEXT:    st.param.b32 [func_retval0+0], %r12;
+; CHECK-NEXT:    ret;
+entry:
+  %v = load <4 x i8>, ptr %p, align 4
+  %cmp = icmp eq <4 x i8> %v, zeroinitializer
+  %sext = sext <4 x i1> %cmp to <4 x i8>
+  ret <4 x i8> %sext
+}