[llvm] [AMDGPU] Skip handling non-byte types in promote alloca. (PR #128769)

[Sumanth Gundapaneni via llvm-commits]

https://github.com/sgundapa created https://github.com/llvm/llvm-project/pull/128769

Non-byte types like i1 can be packed and be supported. For the time being these types are not promoted.

Issue found by fuzzer.

>From 4082aa0a48bc20816082bba8ebf0b5b4590b97ad Mon Sep 17 00:00:00 2001
From: Sumanth Gundapaneni <sumanth.gundapaneni at amd.com>
Date: Tue, 25 Feb 2025 14:18:08 -0600
Subject: [PATCH] [AMDGPU] Skip handling non-byte types in promote alloca.

Non-byte types like i1 can be packed and be supported. For the time being
these types are not promoted.

Issue found by fuzzer.
---
 .../lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp | 11 ++++++++--
 .../promote-alloca-skip-non-byte-type.ll      | 21 +++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/AMDGPU/promote-alloca-skip-non-byte-type.ll

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp b/llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp
index 28016b5936ccf..007f930cea4f3 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp
@@ -759,6 +759,14 @@ bool AMDGPUPromoteAllocaImpl::tryPromoteAllocaToVector(AllocaInst &Alloca) {
     return false;
   }
 
+  Type *VecEltTy = VectorTy->getElementType();
+  constexpr unsigned SIZE_OF_BYTE = 8;
+  unsigned ElementSizeInBits = DL->getTypeSizeInBits(VecEltTy);
+  // FIXME: The non-byte type like i1 can be packed and be supported, but
+  // currently we do not handle them.
+  if (ElementSizeInBits % SIZE_OF_BYTE != 0)
+    return false;
+
   std::map<GetElementPtrInst *, WeakTrackingVH> GEPVectorIdx;
   SmallVector<Instruction *> WorkList;
   SmallVector<Instruction *> UsersToRemove;
@@ -776,8 +784,7 @@ bool AMDGPUPromoteAllocaImpl::tryPromoteAllocaToVector(AllocaInst &Alloca) {
 
   LLVM_DEBUG(dbgs() << "  Attempting promotion to: " << *VectorTy << "\n");
 
-  Type *VecEltTy = VectorTy->getElementType();
-  unsigned ElementSize = DL->getTypeSizeInBits(VecEltTy) / 8;
+  unsigned ElementSize = ElementSizeInBits / SIZE_OF_BYTE;
   for (auto *U : Uses) {
     Instruction *Inst = cast<Instruction>(U->getUser());
 
diff --git a/llvm/test/CodeGen/AMDGPU/promote-alloca-skip-non-byte-type.ll b/llvm/test/CodeGen/AMDGPU/promote-alloca-skip-non-byte-type.ll
new file mode 100644
index 0000000000000..3d2234f0a7ac3
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/promote-alloca-skip-non-byte-type.ll
@@ -0,0 +1,21 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -mtriple=amdgcn-unknown-amdhsa -passes=amdgpu-promote-alloca < %s | FileCheck %s
+
+; Verify that we do not crash and not promote non-byte alloca types.
+define <8 x i1> @non_byte_alloca_type() {
+; CHECK-LABEL: define <8 x i1> @non_byte_alloca_type() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[C:%.*]] = icmp ugt <16 x i1> zeroinitializer, zeroinitializer
+; CHECK-NEXT:    [[RP:%.*]] = alloca <8 x i1>, align 1
+; CHECK-NEXT:    [[TMP0:%.*]] = load <8 x i1>, ptr [[RP]], align 1
+; CHECK-NEXT:    store <16 x i1> [[C]], ptr [[RP]], align 2
+; CHECK-NEXT:    ret <8 x i1> [[TMP0]]
+;
+entry:
+  %C = icmp ugt <16 x i1> zeroinitializer, zeroinitializer
+  %RP = alloca <8 x i1>, align 1
+  %0 = load <8 x i1>, ptr %RP, align 1
+  store <16 x i1> %C, ptr %RP, align 2
+  ret <8 x i1> %0
+}
+