[SROA] Allow SROA on pointers with invariant group intrinsic uses

[aeubanks]

Summary

When we are able to SROA an alloca, we know all uses of it, meaning we
don't have to preserve the invariant group intrinsics and metadata.

It's possible that we could lose information regarding redundant
loads/stores, but that's unlikely to have any real impact since right
now the only user is Clang and vtables.

[rnk]

lgtm

Part of me thinks we really ought to add an integration test to Clang for this. However, I found it difficult to construct a program that SROA fails to optimize that the follow-up MemorySSA passes can't optimize without SROA, so maybe it would be too fragile.

llvm/include/llvm/IR/Instruction.h

default:
       return false;
     }
   }
 
   /// Return true if the instruction is a llvm.lifetime.start or
   /// llvm.lifetime.end marker.
   bool isLifetimeStartOrEnd() const;
 
+  /// Return true if the instruction is a llvm.launder.invariant.group or
+  /// llvm.strip.invariant.group.
+  bool isLaunderOrStripInvariantGroup() const;
+
   /// Return true if the instruction is a DbgInfoIntrinsic or PseudoProbeInst.
   bool isDebugOrPseudoInst() const;
 
   /// Return a pointer to the next non-debug instruction in the same basic
   /// block as 'this', or nullptr if no such instruction exists. Skip any pseudo
   /// operations if \c SkipPseudoOp is true.
   const Instruction *
   getNextNonDebugInstruction(bool SkipPseudoOp = false) const;

llvm/lib/IR/Instruction.cpp

 //===-- Instruction.cpp - Implement the Instruction class -----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the Instruction class for the IR library.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/IR/Instruction.h"
-#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/Type.h"
 using namespace llvm;
 
 Instruction::Instruction(Type *ty, unsigned it, Use *Ops, unsigned NumOps,
                          Instruction *InsertBefore)
   : User(ty, Value::InstructionVal + it, Ops, NumOps), Parent(nullptr) {
if (const auto *CB = dyn_cast<CallBase>(this))
     // return. Remove this workaround once all intrinsics are appropriately
     // annotated.
     return CB->hasFnAttr(Attribute::WillReturn) ||
            (isa<IntrinsicInst>(CB) && CB->onlyReadsMemory());
   return true;
 }
 
 bool Instruction::isLifetimeStartOrEnd() const {
-  auto II = dyn_cast<IntrinsicInst>(this);
+  auto *II = dyn_cast<IntrinsicInst>(this);
   if (!II)
     return false;
   Intrinsic::ID ID = II->getIntrinsicID();
   return ID == Intrinsic::lifetime_start || ID == Intrinsic::lifetime_end;
 }
 
+bool Instruction::isLaunderOrStripInvariantGroup() const {
+  auto *II = dyn_cast<IntrinsicInst>(this);
+  if (!II)
+    return false;
+  Intrinsic::ID ID = II->getIntrinsicID();
+  return ID == Intrinsic::launder_invariant_group ||
+         ID == Intrinsic::strip_invariant_group;
+}
+
 bool Instruction::isDebugOrPseudoInst() const {
   return isa<DbgInfoIntrinsic>(this) || isa<PseudoProbeInst>(this);
 }
 
 const Instruction *
 Instruction::getNextNonDebugInstruction(bool SkipPseudoOp) const {
   for (const Instruction *I = getNextNode(); I; I = I->getNextNode())
     if (!isa<DbgInfoIntrinsic>(I) && !(SkipPseudoOp && isa<PseudoProbeInst>(I)))

llvm/lib/Transforms/Scalar/SROA.cpp

void visitMemTransferInst(MemTransferInst &II) {
     // Insert the use now that we've fixed up the splittable nature.
     insertUse(II, Offset, Size, /*IsSplittable=*/Inserted && Length);
 
     // Check that we ended up with a valid index in the map.
     assert(AS.Slices[PrevIdx].getUse()->getUser() == &II &&
            "Map index doesn't point back to a slice with this user.");
   }
 
-  // Disable SRoA for any intrinsics except for lifetime invariants.
+  // Disable SRoA for any intrinsics except for lifetime invariants and
+  // invariant group.
   // FIXME: What about debug intrinsics? This matches old behavior, but
   // doesn't make sense.
   void visitIntrinsicInst(IntrinsicInst &II) {
     if (II.isDroppable()) {
       AS.DeadUseIfPromotable.push_back(U);
       return;
     }
 
     if (!IsOffsetKnown)
       return PI.setAborted(&II);
 
     if (II.isLifetimeStartOrEnd()) {
       ConstantInt *Length = cast<ConstantInt>(II.getArgOperand(0));
       uint64_t Size = std::min(AllocSize - Offset.getLimitedValue(),
                                Length->getLimitedValue());
       insertUse(II, Offset, Size, true);
       return;
     }
 
+    if (II.isLaunderOrStripInvariantGroup()) {
+      enqueueUsers(II);
+      return;
+    }
+
     Base::visitIntrinsicInst(II);
   }
 
   Instruction *hasUnsafePHIOrSelectUse(Instruction *Root, uint64_t &Size) {
     // We consider any PHI or select that results in a direct load or store of
     // the same offset to be a viable use for slicing purposes. These uses
     // are considered unsplittable and the size is the maximum loaded or stored
     // size.

llvm/test/Transforms/SROA/invariant-group.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=sroa -S -o - < %s | FileCheck %s
+
+%t = type { i32, i32 }
+
+declare i8* @llvm.launder.invariant.group.p0i8(i8* %a)
+declare i8* @llvm.strip.invariant.group.p0i8(i8* %a)
+declare void @h(i32 %a)
+declare i32 @somevalue()
+
+define void @f() {
+; CHECK-LABEL: @f(
+; CHECK-NEXT:    [[SV1:%.*]] = call i32 @somevalue()
+; CHECK-NEXT:    [[SV2:%.*]] = call i32 @somevalue()
+; CHECK-NEXT:    call void @h(i32 [[SV1]])
+; CHECK-NEXT:    call void @h(i32 [[SV2]])
+; CHECK-NEXT:    ret void
+;
+  %a = alloca %t
+
+  %a1 = getelementptr inbounds %t, %t* %a, i32 0, i32 0
+  %a1_i8 = bitcast i32* %a1 to i8*
+  %a1_i8_inv = call i8* @llvm.launder.invariant.group.p0i8(i8* %a1_i8)
+  %a1_inv = bitcast i8* %a1_i8_inv to i32*
+  %a2 = getelementptr inbounds %t, %t* %a, i32 0, i32 1
+
+  %sv1 = call i32 @somevalue()
+  %sv2 = call i32 @somevalue()
+
+  store i32 %sv1, i32* %a1_inv, !invariant.group !0
+  store i32 %sv2, i32* %a2
+
+  %v1 = load i32, i32* %a1_inv, !invariant.group !0
+  %v2 = load i32, i32* %a2
+
+  call void @h(i32 %v1)
+  call void @h(i32 %v2)
+
+  ret void
+}
+
+define void @g() {
+; CHECK-LABEL: @g(
+; CHECK-NEXT:    [[A:%.*]] = alloca [[T:%.*]], align 8
+; CHECK-NEXT:    [[A1:%.*]] = getelementptr inbounds [[T]], %t* [[A]], i32 0, i32 0
+; CHECK-NEXT:    [[A1_I8:%.*]] = bitcast i32* [[A1]] to i8*
+; CHECK-NEXT:    [[A1_I8_INV:%.*]] = call i8* @llvm.launder.invariant.group.p0i8(i8* [[A1_I8]])
+; CHECK-NEXT:    [[A1_INV:%.*]] = bitcast i8* [[A1_I8_INV]] to i32*
+; CHECK-NEXT:    [[A2:%.*]] = getelementptr inbounds [[T]], %t* [[A]], i32 0, i32 1
+; CHECK-NEXT:    [[SV1:%.*]] = call i32 @somevalue()
+; CHECK-NEXT:    [[SV2:%.*]] = call i32 @somevalue()
+; CHECK-NEXT:    store i32 [[SV1]], i32* [[A1_INV]], align 4, !invariant.group !0
+; CHECK-NEXT:    store i32 [[SV2]], i32* [[A2]], align 4
+; CHECK-NEXT:    [[V1:%.*]] = load i32, i32* [[A1_INV]], align 4, !invariant.group !0
+; CHECK-NEXT:    [[V2:%.*]] = load i32, i32* [[A2]], align 4
+; CHECK-NEXT:    call void @h(i32 [[V1]])
+; CHECK-NEXT:    call void @h(i32 [[V2]])
+; CHECK-NEXT:    [[A1_STRIPPED:%.*]] = call i8* @llvm.strip.invariant.group.p0i8(i8* [[A1_I8]])
+; CHECK-NEXT:    [[A1_INT:%.*]] = ptrtoint i8* [[A1_STRIPPED]] to i32
+; CHECK-NEXT:    call void @h(i32 [[A1_INT]])
+; CHECK-NEXT:    ret void
+;
+  %a = alloca %t
+
+  %a1 = getelementptr inbounds %t, %t* %a, i32 0, i32 0
+  %a1_i8 = bitcast i32* %a1 to i8*
+  %a1_i8_inv = call i8* @llvm.launder.invariant.group.p0i8(i8* %a1_i8)
+  %a1_inv = bitcast i8* %a1_i8_inv to i32*
+  %a2 = getelementptr inbounds %t, %t* %a, i32 0, i32 1
+
+  %sv1 = call i32 @somevalue()
+  %sv2 = call i32 @somevalue()
+
+  store i32 %sv1, i32* %a1_inv, !invariant.group !0
+  store i32 %sv2, i32* %a2
+
+  %v1 = load i32, i32* %a1_inv, !invariant.group !0
+  %v2 = load i32, i32* %a2
+
+  call void @h(i32 %v1)
+  call void @h(i32 %v2)
+
+  %a1_stripped = call i8* @llvm.strip.invariant.group.p0i8(i8* %a1_i8)
+  %a1_int = ptrtoint i8* %a1_stripped to i32
+  call void @h(i32 %a1_int)
+
+  ret void
+}
+
+!0 = !{}