[multi-gpu] Phase 4: air-symmetric-alloc-to-mgpu lowering pass

erwei-xilinx · claude · erwei-xilinx · commit 087fcb789453 · 2026-05-06T04:39:20.000Z
New conversion pass that replaces `memref.alloc` carrying the unit
attribute `air.symmetric` with a call to `mgpuSymmetricAlloc(size, stream)`.
The returned `!llvm.ptr` is wrapped in an LLVM memref descriptor (struct)
and projected back to the original memref type via
`builtin.unrealized_conversion_cast` so downstream uses keep working
through the standard `convert-to-llvm` pipeline.

`memref.dealloc` ops whose operand traces back (through the cast) to a
symmetric alloc are rewritten to `mgpuSymmetricFree`.

The pass is a no-op when no `air.symmetric` allocations are present.

## Files

- `mlir/include/air/Conversion/AIRSymmetricAllocToMgpuPass.h` — header
- `mlir/include/air/Conversion/GPUPasses.td` — `air-symmetric-alloc-to-mgpu` def
- `mlir/include/air/Conversion/GPUPassDetail.h` — `GEN_PASS_DEF_AIRSYMMETRICALLOCTOMGPU`
- `mlir/lib/Conversion/AIRSymmetricAllocToMgpuPass.cpp` — implementation
- `mlir/lib/Conversion/{CMakeLists.txt,Passes.cpp}` — registration
- `mlir/test/Conversion/AIRSymmetricAllocToMgpu/symmetric_alloc.mlir` — FileCheck
- `test/gpu/symmetric_heap_dma/air_sym_with_alloc.mlir` — high-level e2e
  using `memref.alloc {air.symmetric}` (Phase 3 + Phase 4 chained)
- `test/gpu/symmetric_heap_dma/run.sh` — `INPUT=alloc` selector

## Test plan

FileCheck unit tests:
- 1D alloc + dealloc shape (size, descriptor, cast, free)
- 2D alloc with row-major strides in descriptor
- Element type byte-size: f32 (4B), f64 (8B), i32 (4B)
- Multiple symmetric allocs share one decl pair
- Pass is a no-op for non-symmetric allocs
- Pass is a no-op when there are zero symmetric allocs

End-to-end on rad-mi300a-sh5-1 (SHARE_GPU=1, 2 ranks):
- INPUT=handwritten — PASS (Phase 2 baseline)
- INPUT=rank — PASS (Phase 3)
- INPUT=alloc — PASS (Phase 4: chained Phase 4 + Phase 3 lowering)

Co-Authored-By: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/mlir/include/air/Conversion/AIRSymmetricAllocToMgpuPass.h b/mlir/include/air/Conversion/AIRSymmetricAllocToMgpuPass.h
@@ -0,0 +1,22 @@
+//===- AIRSymmetricAllocToMgpuPass.h ----------------------------*- C++ -*-===//
+//
+// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+//===-----------------------------------------------------------------------===//
+
+#ifndef AIR_CONVERSION_AIR_SYMMETRIC_ALLOC_TO_MGPU_PASS_H
+#define AIR_CONVERSION_AIR_SYMMETRIC_ALLOC_TO_MGPU_PASS_H
+
+#include "mlir/Pass/Pass.h"
+#include <memory>
+
+namespace xilinx {
+namespace air {
+
+std::unique_ptr<mlir::Pass> createAIRSymmetricAllocToMgpuPass();
+
+} // namespace air
+} // namespace xilinx
+
+#endif // AIR_CONVERSION_AIR_SYMMETRIC_ALLOC_TO_MGPU_PASS_H
diff --git a/mlir/include/air/Conversion/GPUPassDetail.h b/mlir/include/air/Conversion/GPUPassDetail.h
@@ -27,6 +27,7 @@ using namespace mlir;
 #define GEN_PASS_DEF_CONVERTAIRTOROCDL
 #define GEN_PASS_DEF_CONVERTGPUKERNELOUTLINE
 #define GEN_PASS_DEF_AIRRANKTOMGPU
+#define GEN_PASS_DEF_AIRSYMMETRICALLOCTOMGPU
 #include "air/Conversion/GPUPasses.h.inc"
 
 } // namespace air
diff --git a/mlir/include/air/Conversion/GPUPasses.td b/mlir/include/air/Conversion/GPUPasses.td
@@ -49,6 +49,29 @@ def ConvertGPUKernelOutline : Pass<"air-gpu-outlining", "ModuleOp"> {
   let options = [];
 }
 
+def AIRSymmetricAllocToMgpu : Pass<"air-symmetric-alloc-to-mgpu", "ModuleOp"> {
+  let summary = "Lower memref.alloc {air.symmetric} to mgpuSymmetricAlloc and "
+                "memref.dealloc of the result to mgpuSymmetricFree";
+  let constructor = "xilinx::air::createAIRSymmetricAllocToMgpuPass()";
+  let description = [{
+    Replaces each `memref.alloc` carrying the unit attribute `air.symmetric`
+    with a call to `mgpuSymmetricAlloc(size_in_bytes, stream)` returning
+    `!llvm.ptr`, then builds an LLVM memref descriptor (struct) wrapping that
+    pointer and projects it back to the original memref type via
+    `builtin.unrealized_conversion_cast` so downstream uses keep working.
+
+    For every `memref.dealloc` whose operand traces back (through a single
+    `unrealized_conversion_cast`) to such a symmetric alloc, the pass emits
+    `mgpuSymmetricFree(ptr, stream)` and erases the dealloc.
+
+    Should run before `convert-to-llvm`. Does nothing if no `air.symmetric`
+    allocations are present.
+  }];
+  let dependentDialects = [
+    "func::FuncDialect", "arith::ArithDialect", "LLVM::LLVMDialect"
+  ];
+}
+
 def AIRRankToMgpu : Pass<"air-rank-to-mgpu", "ModuleOp"> {
   let summary = "Lower air.rank to mgpu* runtime calls (multi-GPU process model)";
   let constructor = "xilinx::air::createAIRRankToMgpuPass()";
diff --git a/mlir/lib/Conversion/AIRSymmetricAllocToMgpuPass.cpp b/mlir/lib/Conversion/AIRSymmetricAllocToMgpuPass.cpp
@@ -0,0 +1,199 @@
+//===- AIRSymmetricAllocToMgpuPass.cpp -------------------------*- C++ -*-===//
+//
+// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+//===-----------------------------------------------------------------------===//
+//
+// Lower memref.alloc carrying the `air.symmetric` attribute to a call to the
+// runtime function `mgpuSymmetricAlloc`. The returned `!llvm.ptr` is wrapped
+// in an LLVM memref descriptor (struct) and projected back to the original
+// memref type via `builtin.unrealized_conversion_cast` so that downstream
+// uses keep working.
+//
+// `memref.dealloc` ops whose operand traces (through a single
+// `unrealized_conversion_cast`) back to a symmetric alloc are rewritten to
+// `mgpuSymmetricFree`.
+//
+//===-----------------------------------------------------------------------===//
+
+#include "air/Conversion/AIRSymmetricAllocToMgpuPass.h"
+#include "air/Conversion/GPUPassDetail.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+using namespace xilinx;
+
+namespace {
+
+// Ensure a private extern func declaration exists at module scope.
+static func::FuncOp ensureExternFunc(ModuleOp module, OpBuilder &builder,
+                                     StringRef name, FunctionType type) {
+  if (auto fn = module.lookupSymbol<func::FuncOp>(name))
+    return fn;
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPointToStart(module.getBody());
+  auto fn = func::FuncOp::create(builder, module.getLoc(), name, type);
+  fn.setPrivate();
+  return fn;
+}
+
+// Compute the byte size of a static-shaped memref as an i64 SSA value.
+// Returns nullptr if the memref is dynamically shaped.
+static Value computeMemrefByteSize(OpBuilder &b, Location loc, MemRefType ty) {
+  if (!ty.hasStaticShape())
+    return nullptr;
+  int64_t numElts = 1;
+  for (int64_t d : ty.getShape())
+    numElts *= d;
+  unsigned eltBits = ty.getElementType().getIntOrFloatBitWidth();
+  if (eltBits == 0 || (eltBits % 8) != 0)
+    return nullptr;
+  int64_t totalBytes = numElts * (eltBits / 8);
+  return arith::ConstantOp::create(b, loc, b.getI64Type(),
+                                   b.getI64IntegerAttr(totalBytes));
+}
+
+// Build an LLVM memref descriptor struct populated with the given pointer.
+// For now we support only static-shape, contiguous, identity-layout memrefs
+// without an offset. For dimensions: sizes from the type, strides as
+// row-major (innermost stride = 1).
+static Value buildMemrefDescriptor(OpBuilder &b, Location loc,
+                                   MemRefType memrefTy, Value ptr) {
+  ArrayRef<int64_t> shape = memrefTy.getShape();
+  unsigned rank = shape.size();
+  auto i64Ty = b.getI64Type();
+  auto ptrTy = LLVM::LLVMPointerType::get(b.getContext());
+
+  // Build the descriptor type: !llvm.struct<(ptr, ptr, i64, array<R x i64>,
+  // array<R x i64>)>. For rank-0 memrefs, MLIR omits the size/stride arrays.
+  SmallVector<Type, 5> descFields;
+  descFields.push_back(ptrTy);
+  descFields.push_back(ptrTy);
+  descFields.push_back(i64Ty);
+  if (rank > 0) {
+    descFields.push_back(LLVM::LLVMArrayType::get(i64Ty, rank));
+    descFields.push_back(LLVM::LLVMArrayType::get(i64Ty, rank));
+  }
+  auto structTy = LLVM::LLVMStructType::getLiteral(b.getContext(), descFields);
+
+  Value desc = LLVM::PoisonOp::create(b, loc, structTy);
+  desc = LLVM::InsertValueOp::create(b, loc, desc, ptr, ArrayRef<int64_t>{0});
+  desc = LLVM::InsertValueOp::create(b, loc, desc, ptr, ArrayRef<int64_t>{1});
+  Value zero = LLVM::ConstantOp::create(b, loc, i64Ty, b.getI64IntegerAttr(0));
+  desc = LLVM::InsertValueOp::create(b, loc, desc, zero, ArrayRef<int64_t>{2});
+
+  if (rank > 0) {
+    // Compute row-major strides from shape (innermost = 1).
+    SmallVector<int64_t> strides(rank, 1);
+    for (int i = static_cast<int>(rank) - 2; i >= 0; --i)
+      strides[i] = strides[i + 1] * shape[i + 1];
+    for (unsigned i = 0; i < rank; ++i) {
+      Value sz = LLVM::ConstantOp::create(b, loc, i64Ty,
+                                          b.getI64IntegerAttr(shape[i]));
+      desc = LLVM::InsertValueOp::create(b, loc, desc, sz,
+                                         ArrayRef<int64_t>{3, (int64_t)i});
+      Value st = LLVM::ConstantOp::create(b, loc, i64Ty,
+                                          b.getI64IntegerAttr(strides[i]));
+      desc = LLVM::InsertValueOp::create(b, loc, desc, st,
+                                         ArrayRef<int64_t>{4, (int64_t)i});
+    }
+  }
+  return desc;
+}
+
+struct AIRSymmetricAllocToMgpuPass
+    : public xilinx::air::impl::AIRSymmetricAllocToMgpuBase<
+          AIRSymmetricAllocToMgpuPass> {
+
+  AIRSymmetricAllocToMgpuPass() = default;
+  AIRSymmetricAllocToMgpuPass(const AIRSymmetricAllocToMgpuPass &) {}
+
+  void runOnOperation() override {
+    auto module = getOperation();
+    OpBuilder builder(module.getContext());
+    auto i64Ty = builder.getI64Type();
+    auto ptrTy = LLVM::LLVMPointerType::get(module.getContext());
+
+    // Collect symmetric allocs.
+    SmallVector<memref::AllocOp> symAllocs;
+    module.walk([&](memref::AllocOp op) {
+      if (op->hasAttr("air.symmetric"))
+        symAllocs.push_back(op);
+    });
+
+    if (symAllocs.empty())
+      return;
+
+    auto allocFn = ensureExternFunc(
+        module, builder, "mgpuSymmetricAlloc",
+        builder.getFunctionType({i64Ty, ptrTy}, {ptrTy}));
+    auto freeFn = ensureExternFunc(
+        module, builder, "mgpuSymmetricFree",
+        builder.getFunctionType({ptrTy, ptrTy}, {}));
+
+    // Track the !llvm.ptr backing each lowered memref so deallocs can look
+    // them up.
+    DenseMap<Value, Value> symmetricMemrefToPtr;
+
+    for (memref::AllocOp alloc : symAllocs) {
+      auto memrefTy = alloc.getType();
+      Location loc = alloc.getLoc();
+      builder.setInsertionPoint(alloc);
+
+      Value sizeBytes = computeMemrefByteSize(builder, loc, memrefTy);
+      if (!sizeBytes) {
+        alloc.emitOpError(
+            "air.symmetric memref.alloc requires a static-shape memref with "
+            "byte-aligned element type");
+        signalPassFailure();
+        return;
+      }
+      Value nullPtr = LLVM::ZeroOp::create(builder, loc, ptrTy);
+      Value ptr = func::CallOp::create(builder, loc, allocFn,
+                                        ValueRange{sizeBytes, nullPtr})
+                       .getResult(0);
+
+      Value desc = buildMemrefDescriptor(builder, loc, memrefTy, ptr);
+      Value newMemref = UnrealizedConversionCastOp::create(
+                            builder, loc, TypeRange{memrefTy}, ValueRange{desc})
+                            .getResult(0);
+      symmetricMemrefToPtr[newMemref] = ptr;
+      alloc.getResult().replaceAllUsesWith(newMemref);
+      alloc.erase();
+    }
+
+    // Lower deallocs whose operand traces back to a symmetric alloc.
+    SmallVector<memref::DeallocOp> deallocs;
+    module.walk([&](memref::DeallocOp op) { deallocs.push_back(op); });
+    for (memref::DeallocOp d : deallocs) {
+      Value src = d.getMemref();
+      auto it = symmetricMemrefToPtr.find(src);
+      if (it == symmetricMemrefToPtr.end())
+        continue; // not a symmetric memref
+      builder.setInsertionPoint(d);
+      Value nullPtr = LLVM::ZeroOp::create(builder, d.getLoc(), ptrTy);
+      func::CallOp::create(builder, d.getLoc(), freeFn,
+                            ValueRange{it->second, nullPtr});
+      d.erase();
+    }
+  }
+};
+
+} // namespace
+
+namespace xilinx {
+namespace air {
+
+std::unique_ptr<mlir::Pass> createAIRSymmetricAllocToMgpuPass() {
+  return std::make_unique<AIRSymmetricAllocToMgpuPass>();
+}
+
+} // namespace air
+} // namespace xilinx
diff --git a/mlir/lib/Conversion/CMakeLists.txt b/mlir/lib/Conversion/CMakeLists.txt
@@ -58,6 +58,7 @@ if(AIR_ENABLE_GPU)
     AIRTranslateToLLVMPass.cpp
     GPUKernelOutlinePass.cpp
     AIRRankToMgpuPass.cpp
+    AIRSymmetricAllocToMgpuPass.cpp
   )
   list(APPEND CONVERSION_LINK_LIBS
     MLIRGPUDialect
diff --git a/mlir/lib/Conversion/Passes.cpp b/mlir/lib/Conversion/Passes.cpp
@@ -10,6 +10,7 @@
 
 #if AIR_ENABLE_GPU
 #include "air/Conversion/AIRRankToMgpuPass.h"
+#include "air/Conversion/AIRSymmetricAllocToMgpuPass.h"
 #include "air/Conversion/AIRToROCDLPass.h"
 #include "air/Conversion/AIRTranslateToLLVMPass.h"
 #include "air/Conversion/GPUKernelOutlinePass.h"
diff --git a/mlir/test/Conversion/AIRSymmetricAllocToMgpu/symmetric_alloc.mlir b/mlir/test/Conversion/AIRSymmetricAllocToMgpu/symmetric_alloc.mlir
@@ -0,0 +1,106 @@
+//===- symmetric_alloc.mlir -------------------------------------*- MLIR -*-===//
+//
+// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+//===-----------------------------------------------------------------------===//
+
+// RUN: air-opt %s --split-input-file -air-symmetric-alloc-to-mgpu | FileCheck %s
+
+// Basic 1D alloc + dealloc.
+// CHECK-LABEL: func.func @basic_alloc_dealloc
+// CHECK: %[[SZ:.*]] = arith.constant 4096 : i64
+// CHECK: %[[NULL:.*]] = llvm.mlir.zero : !llvm.ptr
+// CHECK: %[[PTR:.*]] = call @mgpuSymmetricAlloc(%[[SZ]], %[[NULL]]) : (i64, !llvm.ptr) -> !llvm.ptr
+// Descriptor build (poison + insertvalue) then unrealized cast.
+// CHECK: llvm.mlir.poison
+// CHECK: llvm.insertvalue %[[PTR]]
+// CHECK: llvm.insertvalue %[[PTR]]
+// CHECK: builtin.unrealized_conversion_cast {{.*}} : !llvm.struct<{{.*}}> to memref<1024xf32>
+// Dealloc -> mgpuSymmetricFree.
+// CHECK: call @mgpuSymmetricFree(%[[PTR]],
+// CHECK-NOT: memref.alloc
+// CHECK-NOT: memref.dealloc
+func.func @basic_alloc_dealloc() {
+  %buf = memref.alloc() {air.symmetric} : memref<1024xf32>
+  memref.dealloc %buf : memref<1024xf32>
+  return
+}
+
+// -----
+
+// 2D alloc: 64*64*4 = 16384 bytes; descriptor strides should be [64, 1].
+// CHECK-LABEL: func.func @alloc_2d
+// CHECK: arith.constant 16384 : i64
+// CHECK: call @mgpuSymmetricAlloc
+// Strides 64 then 1 in the descriptor (innermost-most-contiguous).
+// CHECK: llvm.mlir.constant(64 : i64)
+// CHECK: llvm.insertvalue
+// CHECK: llvm.mlir.constant(1 : i64)
+// CHECK: llvm.insertvalue
+// CHECK: builtin.unrealized_conversion_cast {{.*}} : !llvm.struct<{{.*}}> to memref<64x64xf32>
+func.func @alloc_2d() -> memref<64x64xf32> {
+  %buf = memref.alloc() {air.symmetric} : memref<64x64xf32>
+  return %buf : memref<64x64xf32>
+}
+
+// -----
+
+// f64 element type (8 bytes): 1024 * 8 = 8192 bytes.
+// CHECK-LABEL: func.func @f64_element
+// CHECK: arith.constant 8192 : i64
+func.func @f64_element() {
+  %buf = memref.alloc() {air.symmetric} : memref<1024xf64>
+  memref.dealloc %buf : memref<1024xf64>
+  return
+}
+
+// -----
+
+// i32 element type (4 bytes): 256 * 4 = 1024 bytes.
+// CHECK-LABEL: func.func @i32_element
+// CHECK: arith.constant 1024 : i64
+func.func @i32_element() {
+  %buf = memref.alloc() {air.symmetric} : memref<256xi32>
+  memref.dealloc %buf : memref<256xi32>
+  return
+}
+
+// -----
+
+// Multiple symmetric allocs in one function: each lowered independently;
+// extern decls are emitted exactly once at module scope.
+// Match the actual emission order: Free decl before Alloc decl.
+// CHECK-COUNT-1: func.func private @mgpuSymmetricFree
+// CHECK-NOT: func.func private @mgpuSymmetricFree
+// CHECK-COUNT-1: func.func private @mgpuSymmetricAlloc
+// CHECK-NOT: func.func private @mgpuSymmetricAlloc
+// CHECK-LABEL: func.func @two_allocs
+// CHECK-COUNT-2: call @mgpuSymmetricAlloc
+// CHECK-COUNT-2: call @mgpuSymmetricFree
+func.func @two_allocs() {
+  %a = memref.alloc() {air.symmetric} : memref<32xf32>
+  %b = memref.alloc() {air.symmetric} : memref<64xf32>
+  memref.dealloc %a : memref<32xf32>
+  memref.dealloc %b : memref<64xf32>
+  return
+}
+
+// -----
+
+// LAST partition: cases that test the pass leaves things untouched.
+// Both `ignores_non_symmetric` and `no_symmetric_alloc` are folded here
+// so the trailing CHECK-NOTs only need to match against this one (final)
+// partition's text.
+// CHECK-LABEL: func.func @no_symmetric_changes
+// CHECK: memref.alloc() : memref<1024xf32>
+// CHECK: memref.alloc() : memref<32xf32>
+// CHECK-NOT: mgpuSymmetricAlloc
+// CHECK-NOT: mgpuSymmetricFree
+func.func @no_symmetric_changes() {
+  %a = memref.alloc() : memref<1024xf32>
+  memref.dealloc %a : memref<1024xf32>
+  %b = memref.alloc() : memref<32xf32>
+  memref.dealloc %b : memref<32xf32>
+  return
+}
diff --git a/test/gpu/symmetric_heap_dma/air_sym_with_alloc.mlir b/test/gpu/symmetric_heap_dma/air_sym_with_alloc.mlir
diff --git a/test/gpu/symmetric_heap_dma/run.sh b/test/gpu/symmetric_heap_dma/run.sh

Original file line number	Diff line number	Diff line change
`@@ -58,6 +58,7 @@ if(AIR_ENABLE_GPU)`
`58`	`58`	`AIRTranslateToLLVMPass.cpp`
`59`	`59`	`GPUKernelOutlinePass.cpp`
`60`	`60`	`AIRRankToMgpuPass.cpp`
	`61`	`+ AIRSymmetricAllocToMgpuPass.cpp`
`61`	`62`	`)`
`62`	`63`	`list(APPEND CONVERSION_LINK_LIBS`
`63`	`64`	`MLIRGPUDialect`