[llvm] [IR] Add logical SROA pass (PR #192058)

[Antonio Frighetto via llvm-commits]

Nathan =?utf-8?q?Gauër?= <brioche at google.com>,
Nathan =?utf-8?q?Gauër?= <brioche at google.com>,
Nathan =?utf-8?q?Gauër?= <brioche at google.com>,
Nathan =?utf-8?q?Gauër?= <brioche at google.com>,
Nathan =?utf-8?q?Gauër?= <brioche at google.com>,
Nathan =?utf-8?q?Gauër?= <brioche at google.com>,
Nathan =?utf-8?q?Gauër?= <brioche at google.com>
Message-ID:
In-Reply-To: <llvm.org/llvm/llvm-project/pull/192058 at github.com>


================
@@ -0,0 +1,222 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 transformation implements the well known scalar replacement of
+/// aggregates transformation. It tries to identify promotable elements of an
+/// aggregate alloca, and promote them to registers. It will also try to
+/// convert uses of an element (or set of elements) of an alloca into a vector
+/// or bitfield-style integer scalar if appropriate.
+///
+/// It works to do this with minimal slicing of the alloca so that regions
+/// which are merely transferred in and out of external memory remain unchanged
+/// and are not decomposed to scalar code.
+///
+/// Because this also performs alloca promotion, it can be thought of as also
+/// serving the purpose of SSA formation. The algorithm iterates on the
+/// function until all opportunities for promotion have been realized.
+///
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Scalar/LSROA.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/DomTreeUpdater.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Transforms/Scalar.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "logical-sroa"
+
+namespace {
+
+// Return all lifetime intrinsics with the instruction I as operand.
+SmallVector<LifetimeIntrinsic *>
+collectLifetimeIntrinsicsUsing(Instruction &I) {
+  SmallVector<LifetimeIntrinsic *> Output;
+
+  for (const auto &user : I.users()) {
+    auto II = dyn_cast<IntrinsicInst>(user);
+    if (II && isLifetimeIntrinsic(II->getIntrinsicID()))
+      Output.push_back(cast<LifetimeIntrinsic>(II));
+  }
+
+  return Output;
+}
+
+using SGEPVec = SmallVector<StructuredGEPInst *>;
+
+// Returns a vector with one element for each field of the struct allocated by
+// SAI. Each element is a vector of SGEP instruction referencing this field.
+//
+// If any user of SAI is not an SGEP, or an SGEP referencing the whole struct,
+// this function returns an empty array. This function ignores lifetime
+// intrinsics.
+SmallVector<SmallVector<StructuredGEPInst *>>
+collectPerFieldSGEP(StructuredAllocaInst &SAI) {
+  StructType *ST = cast<StructType>(SAI.getAllocationType());
+  SmallVector<SmallVector<StructuredGEPInst *>> Output(ST->getNumElements());
+
+  for (const auto &user : SAI.users()) {
+    auto II = dyn_cast<IntrinsicInst>(user);
+    if (II && II->isLifetimeStartOrEnd())
+      continue;
+
+    auto SGEP = dyn_cast<StructuredGEPInst>(user);
+    if (!SGEP)
+      return {};
+
+    // If the SGEP has no indices, this means we have a pointer on the whole
+    // struct. For now, we bail out: if it was not used, it would be DCE'd, so
+    // there is probably a reference to the whole struct somewhere.
+    if (SGEP->getNumIndices() == 0)
+      return {};
+
+    // IR rule: SGEP on struct can only use constant int as indices.
+    ConstantInt *Index = cast<ConstantInt>(SGEP->getIndexOperand(0));
+    assert(Index->getZExtValue() < Output.size());
+    Output[Index->getZExtValue()].push_back(SGEP);
+  }
+
+  return Output;
+}
+
+// For each lifetime intrinsic in LifetimeIntrinsics, creates a new one, but
+// uses V as operand.
+void copyLifetimeIntrinsicFor(IRBuilder<> &B, LifetimeIntrinsic *II, Value *V) {
+  if (II->getIntrinsicID() == Intrinsic::lifetime_start) {
+    B.SetInsertPoint(II);
+    B.CreateLifetimeStart(V);
+  } else if (II->getIntrinsicID() == Intrinsic::lifetime_end) {
+    B.SetInsertPoint(II);
+    B.CreateLifetimeEnd(V);
+  } else
+    llvm_unreachable("invalid argument: expected a lifetime intrinsic");
+}
+
+void rewriteSGEPChain(IRBuilder<> &B, StructuredGEPInst *SGEP,
+                      StructuredAllocaInst *FieldAlloca) {
+  if (SGEP->getNumIndices() == 1) {
+    SGEP->replaceAllUsesWith(FieldAlloca);
+    SGEP->eraseFromParent();
+    return;
+  }
+
+  SmallVector<Value *, 4> Indices;
+  for (unsigned J = 1; J < SGEP->getNumIndices(); ++J)
+    Indices.push_back(SGEP->getIndexOperand(J));
+
+  B.SetInsertPoint(SGEP);
+  auto *I = B.CreateStructuredGEP(FieldAlloca->getAllocationType(), FieldAlloca,
+                                  Indices, SGEP->getName());
+  SGEP->replaceAllUsesWith(I);
+  SGEP->eraseFromParent();
+}
+
+bool runOnStructuredAlloca(StructuredAllocaInst &SAI) {
+  // For now, LSROA only handles SGEP on structs.
+  StructType *ST = dyn_cast<StructType>(SAI.getAllocationType());
+  if (!ST)
+    return false;
+
+  SmallVector<LifetimeIntrinsic *> LifetimeIntrinsics =
+      collectLifetimeIntrinsicsUsing(SAI);
+  auto PerFieldSGEP = collectPerFieldSGEP(SAI);
+  if (PerFieldSGEP.size() == 0)
+    return false;
+
+  IRBuilder B(&SAI);
+  for (size_t I = 0; I < PerFieldSGEP.size(); ++I) {
+    auto &Users = PerFieldSGEP[I];
+    if (Users.size() == 0)
+      continue;
+
+    B.SetInsertPoint(&SAI);
----------------
antoniofrighetto wrote:

Not sure if it would be better to avoid SetInsertPoint at each iteration by not passing it as argument in the following routines, and instead let them have their own Builder?

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