[llvm] [InstCombine] Skip GEP canonicalization for types containing target extension types (PR #191790)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Maosu Zhao (zhaomaosu)

<details>
<summary>Changes</summary>

GEP canonicalization in InstCombine rewrites single-index GEPs to use an [N x i8] element type. This transformation is invalid when the GEP's element type contains a target extension type, because target extension types carry backend-defined semantics that must be preserved through IR.

---
Full diff: https://github.com/llvm/llvm-project/pull/191790.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstructionCombining.cpp (+20-1) 
- (added) llvm/test/Transforms/InstCombine/gep-target-ext-type.ll (+37) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 6798493de1aa3..ab2ff44b63709 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -3479,13 +3479,32 @@ Instruction *InstCombinerImpl::visitGetElementPtrInst(GetElementPtrInst &GEP) {
         BackIndices, GEP.getNoWrapFlags());
   }
 
+  // Recursively check whether a type contains a target extension type.
+  // Target extension types have backend-defined semantics and must not be
+  // folded or reinterpreted by generic GEP canonicalization.
+  std::function<bool(Type *)> IsTargetExtType = [&](Type *Ty) -> bool {
+    if (isa<TargetExtType>(Ty))
+      return true;
+
+    if (Ty->isVectorTy())
+      return IsTargetExtType(Ty->getScalarType());
+
+    if (Ty->isArrayTy())
+      return IsTargetExtType(Ty->getArrayElementType());
+
+    if (auto *STy = dyn_cast<StructType>(Ty))
+      return any_of(STy->elements(), IsTargetExtType);
+
+    return false;
+  };
   // Canonicalize gep %T to gep [sizeof(%T) x i8]:
   auto IsCanonicalType = [](Type *Ty) {
     if (auto *AT = dyn_cast<ArrayType>(Ty))
       Ty = AT->getElementType();
     return Ty->isIntegerTy(8);
   };
-  if (Indices.size() == 1 && !IsCanonicalType(GEPEltType)) {
+  if (Indices.size() == 1 && !IsCanonicalType(GEPEltType) &&
+      !IsTargetExtType(GEPEltType)) {
     TypeSize Scale = DL.getTypeAllocSize(GEPEltType);
     assert(!Scale.isScalable() && "Should have been handled earlier");
     Type *NewElemTy = Builder.getInt8Ty();
diff --git a/llvm/test/Transforms/InstCombine/gep-target-ext-type.ll b/llvm/test/Transforms/InstCombine/gep-target-ext-type.ll
new file mode 100644
index 0000000000000..49f094bdee159
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/gep-target-ext-type.ll
@@ -0,0 +1,37 @@
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+; Tests that GEP canonicalization (to [N x i8] element type) is suppressed
+; when the element type contains a target extension type.
+
+target datalayout = "e-p:64:64"
+
+; GEP over a nested struct: target extension type is in an inner struct.
+; Exercises the recursive descent through the struct branch.
+define ptr @gep_nested_struct_target_ext(ptr %p, i64 %idx) {
+; CHECK-LABEL: @gep_nested_struct_target_ext(
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr { i32, { target("spirv.DeviceEvent") } }, ptr [[P:%.*]], i64 [[IDX:%.*]]
+; CHECK-NEXT:    ret ptr [[GEP]]
+;
+  %gep = getelementptr { i32, { target("spirv.DeviceEvent") } }, ptr %p, i64 %idx
+  ret ptr %gep
+}
+
+; GEP over an array of target extension types.
+define ptr @gep_array_of_target_ext(ptr %p, i64 %idx) {
+; CHECK-LABEL: @gep_array_of_target_ext(
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [4 x target("spirv.DeviceEvent")], ptr [[P:%.*]], i64 [[IDX:%.*]]
+; CHECK-NEXT:    ret ptr [[GEP]]
+;
+  %gep = getelementptr [4 x target("spirv.DeviceEvent")], ptr %p, i64 %idx
+  ret ptr %gep
+}
+
+; GEP over a direct target extension type (base case).
+define ptr @gep_direct_target_ext(ptr %p, i64 %idx) {
+; CHECK-LABEL: @gep_direct_target_ext(
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr target("spirv.DeviceEvent"), ptr [[P:%.*]], i64 [[IDX:%.*]]
+; CHECK-NEXT:    ret ptr [[GEP]]
+;
+  %gep = getelementptr target("spirv.DeviceEvent"), ptr %p, i64 %idx
+  ret ptr %gep
+}

``````````

</details>


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