[llvm-branch-commits] [mlir] [MLIR][OpenMP] Support allocations of device shared memory (PR #150924)

[Pranav Bhandarkar via llvm-branch-commits]

================
@@ -1102,12 +1102,64 @@ struct DeferredStore {
 };
 } // namespace
 
+/// Check whether allocations for the given operation might potentially have to
+/// be done in device shared memory. That means we're compiling for a offloading
+/// target, the operation is an `omp::TargetOp` or nested inside of one and that
+/// target region represents a Generic (non-SPMD) kernel.
+///
+/// This represents a necessary but not sufficient set of conditions to use
+/// device shared memory in place of regular allocas. Depending on the variable,
+/// its uses or the associated OpenMP construct might also need to be taken into
+/// account.
+static bool
+mightAllocInDeviceSharedMemory(Operation &op,
+                               const llvm::OpenMPIRBuilder &ompBuilder) {
+  if (!ompBuilder.Config.isTargetDevice())
+    return false;
+
+  auto targetOp = dyn_cast<omp::TargetOp>(op);
+  if (!targetOp)
+    targetOp = op.getParentOfType<omp::TargetOp>();
+
+  return targetOp &&
+         !bitEnumContainsAny(
+             targetOp.getKernelExecFlags(targetOp.getInnermostCapturedOmpOp()),
+             omp::TargetRegionFlags::spmd);
+}
+
+/// Check whether the entry block argument representing the private copy of a
+/// variable in an OpenMP construct must be allocated in device shared memory,
+/// based on what the uses of that copy are.
+///
+/// This must only be called if a previous call to
+/// \c mightAllocInDeviceSharedMemory has already returned \c true for the
+/// operation that owns the specified block argument.
+static bool mustAllocPrivateVarInDeviceSharedMemory(BlockArgument value) {
+  Operation *parentOp = value.getOwner()->getParentOp();
+  auto targetOp = dyn_cast<omp::TargetOp>(parentOp);
+  if (!targetOp)
+    targetOp = parentOp->getParentOfType<omp::TargetOp>();
+  assert(targetOp && "expected a parent omp.target operation");
+
+  for (auto *user : value.getUsers()) {
+    if (auto parallelOp = dyn_cast<omp::ParallelOp>(user)) {
+      if (llvm::is_contained(parallelOp.getReductionVars(), value))
+        return true;
+    } else if (auto parallelOp = user->getParentOfType<omp::ParallelOp>()) {
+      if (targetOp->isProperAncestor(parallelOp))
----------------
bhandarkar-pranav wrote:

There seems to be a hole in my understanding of this.  At this point in the code, we know that `value` iis a `BlockArgument`. We know that it has an an ancestor, `targetOp` that is an `omp::TargetOp`. Should't all the users of a `BlockArgument` be such that they are dominated by the `BlockArgument`. Therefore `targetOp` should trivially be an ancestor of all users, no? All this is to say that the use of a `BlockArgument` inside a `omp::ParallelOp` should be enough and this ancestor check is superfluous. Unless, of course, I am missing something that is obvious.

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