[llvm] [SPIRV] Sign-extend operands of sign-sensitive ops on sub-pow2 widths (PR #203661)

[Dmitry Sidorov via llvm-commits]

================
@@ -0,0 +1,234 @@
+; RUN: llc -O0 -verify-machineinstrs -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -verify-machineinstrs -mtriple=spirv64-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+
+; RUN: llc -O0 -verify-machineinstrs -mtriple=spirv32-unknown-unknown %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -verify-machineinstrs -mtriple=spirv32-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+
+; SPIR-V (without sub-byte int extensions) widens sub-pow2 scalars to the next
+; legal width by relabeling the LLT only, without inserting any sign-extension.
+; Sign-sensitive ops (icmp slt/sle/sgt/sge, ashr, sdiv, srem) on such operands
+; would then read the sign bit at the wrong position. The pre-legalizer must
+; emit a sign-extend-in-register before the widening so the wide-width signed
+; op observes the correct sign bit.
+
+; CHECK-DAG: %[[#I8:]] = OpTypeInt 8 0
+; CHECK-DAG: %[[#I32:]] = OpTypeInt 32 0
+; CHECK-DAG: %[[#K4:]] = OpConstant %[[#I8]] 4
+; CHECK-DAG: %[[#K8:]] = OpConstant %[[#I32]] 8
+
+; ----------------------------------------------------------------------------
+; icmp slt i4 against zero (the canonical XLA F4E2M1FN sign-bit-check pattern).
+; CHECK: OpFunction
+; CHECK: %[[#X1:]] = OpFunctionParameter
+; CHECK: OpFunctionParameter
+; CHECK: %[[#SHL1:]] = OpShiftLeftLogical %[[#I8]] %[[#X1]] %[[#K4]]
+; CHECK: %[[#SX1:]] = OpShiftRightArithmetic %[[#I8]] %[[#SHL1]] %[[#K4]]
+; CHECK: OpSLessThan {{%[0-9]+}} %[[#SX1]] {{%[0-9]+}}
+define spir_kernel void @icmp_slt_i4_zero(i4 %x, ptr addrspace(1) %out) {
+  %c = icmp slt i4 %x, 0
+  %r = sext i1 %c to i8
+  store i8 %r, ptr addrspace(1) %out
+  ret void
+}
+
+; ----------------------------------------------------------------------------
+; icmp sle i4: same widening as slt, different SPIR-V opcode.
+; CHECK: OpFunction
+; CHECK: %[[#X_SLE:]] = OpFunctionParameter
+; CHECK: OpFunctionParameter
+; CHECK: %[[#SHL_SLE:]] = OpShiftLeftLogical %[[#I8]] %[[#X_SLE]] %[[#K4]]
+; CHECK: %[[#SX_SLE:]] = OpShiftRightArithmetic %[[#I8]] %[[#SHL_SLE]] %[[#K4]]
+; CHECK: OpSLessThanEqual {{%[0-9]+}} %[[#SX_SLE]] {{%[0-9]+}}
+define spir_kernel void @icmp_sle_i4_zero(i4 %x, ptr addrspace(1) %out) {
+  %c = icmp sle i4 %x, 0
+  %r = sext i1 %c to i8
+  store i8 %r, ptr addrspace(1) %out
+  ret void
+}
+
+; ----------------------------------------------------------------------------
+; icmp sgt i4.
+; CHECK: OpFunction
+; CHECK: %[[#X_SGT:]] = OpFunctionParameter
+; CHECK: OpFunctionParameter
+; CHECK: %[[#SHL_SGT:]] = OpShiftLeftLogical %[[#I8]] %[[#X_SGT]] %[[#K4]]
+; CHECK: %[[#SX_SGT:]] = OpShiftRightArithmetic %[[#I8]] %[[#SHL_SGT]] %[[#K4]]
+; CHECK: OpSGreaterThan {{%[0-9]+}} %[[#SX_SGT]] {{%[0-9]+}}
+define spir_kernel void @icmp_sgt_i4_zero(i4 %x, ptr addrspace(1) %out) {
+  %c = icmp sgt i4 %x, 0
+  %r = sext i1 %c to i8
+  store i8 %r, ptr addrspace(1) %out
+  ret void
+}
+
+; ----------------------------------------------------------------------------
+; icmp sge i4.
+; CHECK: OpFunction
+; CHECK: %[[#X_SGE:]] = OpFunctionParameter
+; CHECK: OpFunctionParameter
+; CHECK: %[[#SHL_SGE:]] = OpShiftLeftLogical %[[#I8]] %[[#X_SGE]] %[[#K4]]
+; CHECK: %[[#SX_SGE:]] = OpShiftRightArithmetic %[[#I8]] %[[#SHL_SGE]] %[[#K4]]
+; CHECK: OpSGreaterThanEqual {{%[0-9]+}} %[[#SX_SGE]] {{%[0-9]+}}
+define spir_kernel void @icmp_sge_i4_zero(i4 %x, ptr addrspace(1) %out) {
+  %c = icmp sge i4 %x, 0
+  %r = sext i1 %c to i8
+  store i8 %r, ptr addrspace(1) %out
+  ret void
+}
+
+; ----------------------------------------------------------------------------
+; icmp slt i4 between two registers: both operands must be sign-extended.
+; CHECK: OpFunction
+; CHECK: %[[#X2:]] = OpFunctionParameter
+; CHECK: %[[#Y2:]] = OpFunctionParameter
+; CHECK: OpFunctionParameter
+; CHECK: %[[#SHLA2:]] = OpShiftLeftLogical %[[#I8]] %[[#X2]] %[[#K4]]
+; CHECK: %[[#SXA2:]] = OpShiftRightArithmetic %[[#I8]] %[[#SHLA2]] %[[#K4]]
+; CHECK: %[[#SHLB2:]] = OpShiftLeftLogical %[[#I8]] %[[#Y2]] %[[#K4]]
+; CHECK: %[[#SXB2:]] = OpShiftRightArithmetic %[[#I8]] %[[#SHLB2]] %[[#K4]]
+; CHECK: OpSLessThan {{%[0-9]+}} %[[#SXA2]] %[[#SXB2]]
+define spir_kernel void @icmp_slt_i4_reg(i4 %x, i4 %y, ptr addrspace(1) %out) {
+  %c = icmp slt i4 %x, %y
+  %r = sext i1 %c to i8
+  store i8 %r, ptr addrspace(1) %out
+  ret void
+}
+
+; ----------------------------------------------------------------------------
+; ashr i4: arithmetic right shift on a widened operand needs the sign bit at
+; the top of the wider register.
+; CHECK: OpFunction
+; CHECK: %[[#X3:]] = OpFunctionParameter
+; CHECK: %[[#Y3:]] = OpFunctionParameter
+; CHECK: OpFunctionParameter
+; CHECK: %[[#SHLA3:]] = OpShiftLeftLogical %[[#I8]] %[[#X3]] %[[#K4]]
+; CHECK: %[[#SXA3:]] = OpShiftRightArithmetic %[[#I8]] %[[#SHLA3]] %[[#K4]]
+; CHECK: %[[#SHLB3:]] = OpShiftLeftLogical %[[#I8]] %[[#Y3]] %[[#K4]]
+; CHECK: %[[#SXB3:]] = OpShiftRightArithmetic %[[#I8]] %[[#SHLB3]] %[[#K4]]
+; CHECK: OpShiftRightArithmetic %[[#I8]] %[[#SXA3]] %[[#SXB3]]
----------------
MrSidims wrote:

This looks sob-optimal, shouldn't it be:
 %shl_x = OpShiftLeftLogical      %i8 %x    %k4
  %sx    = OpShiftRightArithmetic  %i8 %shl_x %k4    ; sign-extend base
  %ym    = OpBitwiseAnd            %i8 %y    %k0x0F  ; mask shift amount
  %res   = OpShiftRightArithmetic  %i8 %sx   %ym
?

https://github.com/llvm/llvm-project/pull/203661