[Mlir-commits] [mlir] [mlir][acc] Add utilities for acc to gpu lowering (PR #209313)

[Slava Zakharin]

================
@@ -206,5 +213,215 @@ void copyParDimsAttr(Operation *from, Operation *to) {
   setParDimsAttr(to, getParDimsAttr(from));
 }
 
+int64_t SharedMemoryBudget::alignOffset(int64_t offset) {
+  return (offset + 15) & ~static_cast<int64_t>(15);
+}
+
+bool SharedMemoryBudget::tryAllocate(int64_t bytes) {
+  int64_t aligned = alignOffset(bytesUsed_);
+  if (aligned + bytes > maxTotalBytes_) {
+    return false;
+  }
+  bytesUsed_ = aligned + bytes;
+  return true;
+}
+
+int64_t sumExistingSharedMemoryBytes(Region &region) {
+  int64_t total = 0;
+  region.walk([&](GPUSharedMemoryOp op) {
+    int64_t upperBound = op.getStaticUpperBoundBytes();
+    total = SharedMemoryBudget::alignOffset(total) + upperBound;
+  });
+  return total;
+}
+
+PrivatizeOp getPrivatizeOp(PrivateLocalOp privateLocal,
+                           ComputeRegionOp computeRegion) {
+  Value value = privateLocal.getPrivatized();
+  if (BlockArgument blockArg = dyn_cast<BlockArgument>(value)) {
+    auto owner = dyn_cast<ComputeRegionOp>(blockArg.getOwner()->getParentOp());
+    value = (owner ? owner : computeRegion).getOperand(blockArg);
+  }
+  PrivatizeOp privatizeOp = value.getDefiningOp<PrivatizeOp>();
+  assert(privatizeOp && "expected privatize op to be the defining op");
+  return privatizeOp;
+}
+
+static bool isThreadXPrivatize(PrivatizeOp privatize) {
+  if (GPUParallelDimsAttr parDimsAttr = privatize.getParDimsAttr())
+    return llvm::any_of(parDimsAttr.getArray(),
+                        [](GPUParallelDimAttr d) { return d.isThreadX(); });
+  return false;
+}
+
+MemRefType getPrivateBaseMemRefType(Type baseTy, ModuleOp module) {
+  auto memrefTy = cast<PointerLikeType>(baseTy).getAsMemRefType(module);
+  assert(memrefTy && "private base type must be convertible to memref");
+  return memrefTy;
+}
+
+SmallVector<GPUParallelDimAttr>
+collectPrivateLocalParDims(PrivateLocalOp privateLocal,
+                           ComputeRegionOp computeRegion) {
+  SmallVector<GPUParallelDimAttr> parDims;
+  auto parentLoop = privateLocal->getParentOfType<scf::ParallelOp>();
+  while (parentLoop) {
+    if (GPUParallelDimsAttr parDimsAttr = getParDimsAttr(parentLoop))
+      for (GPUParallelDimAttr parDim : parDimsAttr.getArray())
+        insertParDim(parDims, parDim);
+    parentLoop = parentLoop->getParentOfType<scf::ParallelOp>();
----------------
vzakhari wrote:

I wonder if we can pass `acc.compute_region` here and start processing its parent `scf.parallel` operations.

It should not be a big deal, as long as the parent loops won't have any par-dims, but we may as well stop at the compute boundary.

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