[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);
----------------
antoniofrighetto wrote:

```suggestion
  auto LifetimeIntrinsics = collectLifetimeIntrinsicsUsing(SAI);
```

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