[llvm] [NVVM][NVPTX] Support spcompress and spdecompress intrinsics (PR #221751)

[Rajat Bajpai via llvm-commits]

================
@@ -0,0 +1,204 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file contains NVVM-specific IR verification logic. These checks are
+/// always compiled and linked as part of LLVMCore.
+///
+//===----------------------------------------------------------------------===//
+
+#include "VerifierInternal.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/IntrinsicsNVPTX.h"
+#include "llvm/Support/MathExtras.h"
+#include <optional>
+
+using namespace llvm;
+
+#define Check(C, ...)                                                          \
+  do {                                                                         \
+    if (!(C)) {                                                                \
+      VS.CheckFailed(__VA_ARGS__);                                             \
+      return;                                                                  \
+    }                                                                          \
+  } while (false)
+
+namespace {
+
+struct SPVectorInfo {
+  unsigned ElemSize;
+  unsigned NumElements;
+  unsigned NumRegisters;
+};
+
+// Register layout of the sparse intrinsic operands, used by the IR verifier.
+struct SPOperandLayout {
+  unsigned MetadataSize;
+  unsigned CompressedDataSize;
+  unsigned DataSize;
+};
+
+} // namespace
+
+// PTX limits the combined vector size of the mdata, cdata, and data operands
+// of spcompress and spdecompress to 253 32-bit registers.
+constexpr unsigned MaxSPOperandRegisters = 253;
+
+static bool isValidSPElemSize(unsigned ElemSize) {
+  return ElemSize == 8 || ElemSize == 16;
+}
+
+static bool isValidSPIdxSize(unsigned IdxSize) {
+  return IdxSize == 2 || IdxSize == 4;
+}
+
+static bool isValidSPRepeatFactor(unsigned RepeatFactor) {
+  return isPowerOf2_32(RepeatFactor) && RepeatFactor <= 64;
+}
+
+static bool isValidSPDecompressFactor(unsigned NumSrc, unsigned NumTgt) {
+  switch (NumSrc) {
+  case 1:
+    return NumTgt == 2 || NumTgt == 4 || NumTgt == 8 || NumTgt == 16;
+  case 2:
+    return NumTgt == 4 || NumTgt == 8 || NumTgt == 16;
+  case 4:
+    return NumTgt == 8 || NumTgt == 16;
+  default:
+    return false;
+  }
+}
+
+static std::optional<SPOperandLayout>
+getSPCompressLayout(unsigned ElemSize, unsigned IdxSize,
+                    unsigned RepeatFactor) {
+  if (!isValidSPElemSize(ElemSize) || !isValidSPIdxSize(IdxSize) ||
+      !isValidSPRepeatFactor(RepeatFactor))
+    return std::nullopt;
+
+  SPOperandLayout Layout = {divideCeil(RepeatFactor * IdxSize, ElemSize),
+                            RepeatFactor, RepeatFactor * 2};
+  if (Layout.MetadataSize + Layout.CompressedDataSize + Layout.DataSize >
+      MaxSPOperandRegisters)
+    return std::nullopt;
+  return Layout;
+}
+
+static std::optional<SPOperandLayout>
+getSPDecompressLayout(unsigned NumSrc, unsigned NumTgt, unsigned ElemSize,
+                      unsigned IdxSize, unsigned RepeatFactor) {
+  if (!isValidSPDecompressFactor(NumSrc, NumTgt) ||
+      !isValidSPElemSize(ElemSize) || !isValidSPIdxSize(IdxSize) ||
+      !isValidSPRepeatFactor(RepeatFactor) || NumSrc * ElemSize > 32 ||
+      (IdxSize == 2 && NumTgt > 4))
+    return std::nullopt;
+
+  unsigned DataBits = NumTgt * ElemSize * RepeatFactor;
+  if (DataBits < 32 || DataBits > 4096)
+    return std::nullopt;
+
+  SPOperandLayout Layout = {divideCeil(NumSrc * IdxSize * RepeatFactor, 32),
+                            divideCeil(NumSrc * ElemSize * RepeatFactor, 32),
+                            divideCeil(DataBits, 32)};
+  if (Layout.MetadataSize + Layout.CompressedDataSize + Layout.DataSize >
+      MaxSPOperandRegisters)
+    return std::nullopt;
+  return Layout;
+}
+
+static std::optional<SPVectorInfo> getSPVectorInfo(Type *Ty) {
+  auto *VT = dyn_cast<FixedVectorType>(Ty);
+  if (!VT)
+    return std::nullopt;
+  auto *ElemTy = dyn_cast<IntegerType>(VT->getElementType());
+  if (!ElemTy)
+    return std::nullopt;
+  unsigned ElemSize = ElemTy->getBitWidth();
+  if (ElemSize != 8 && ElemSize != 16)
+    return std::nullopt;
+  unsigned NumElements = VT->getNumElements();
+  return SPVectorInfo{ElemSize, NumElements,
+                      divideCeil(NumElements, 32 / ElemSize)};
+}
+
+static std::optional<unsigned> getSPMetadataRegisters(Type *Ty) {
+  if (Ty->isIntegerTy(32))
+    return 1;
+  auto *VT = dyn_cast<FixedVectorType>(Ty);
+  if (!VT || !VT->getElementType()->isIntegerTy(32))
+    return std::nullopt;
+  return VT->getNumElements();
----------------
rajatbajpai wrote:

It seems we are supporting <1 x i32> as well, is that intentional? I don't see any test case exercising that.

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