[multi-gpu] Phase 2: air.translate op + air-translate-to-llvm lowering

Introduce an AIR primitive for the symmetric-heap pointer rebase, in
preparation for the kernel-driven producer/consumer redesign per
@mawad-amd's review feedback on PR Xilinx#1577.

  %peer = air.translate %src, %from, %to, %bases : memref<NxT, A>, !llvm.ptr

Signature:
- $source: memref on $from_rank's symmetric heap
- $from_rank, $to_rank: index-typed rank ids
- $heap_bases: !llvm.ptr to the per-rank base table from mgpuGetHeapBases()
- result: same memref type, addressing $to_rank's slice of the same
  collective allocation

The op is Pure and folds when from_rank == to_rank (statically equal SSA
values or matching constant attrs). Naming follows IRIS's `__translate`.

Lowering pass `air-translate-to-llvm` expands each op to the
peer-VA arithmetic plus a freshly-built LLVM memref descriptor:

  byte_diff = ptrtoint(bases[to]) - ptrtoint(bases[from])
  peer_aligned_ptr = src_aligned_ptr + byte_diff   (i8 GEP)
  build descriptor { peer_ptr, peer_ptr, 0, sizes, strides }
  unrealized_conversion_cast back to result memref type

The expansion is pure arithmetic (arith + memref + llvm dialect), no
runtime calls — therefore valid both at host scope and inside `gpu.func`,
provided heap_bases is threaded as a kernel argument.

Tests:
- mlir/test/Dialect/AIR/air_translate.mlir: parser/printer + folder
- mlir/test/Conversion/AIRToROCDL/air_translate_to_llvm.mlir: lowering
  shape on 1D, 2D-addrspace, gpu.func body, and no-op cases

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: erwei-xilinx

mlir/include/air/Conversion/AIRTranslateToLLVMPass.h

@@ -0,0 +1,22 @@
+//===- AIRTranslateToLLVMPass.h ----------------------------------*- C++ -*-===//
+//
+// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+//===-----------------------------------------------------------------------===//
+
+#ifndef AIR_CONVERSION_AIR_TRANSLATE_TO_LLVM_PASS_H
+#define AIR_CONVERSION_AIR_TRANSLATE_TO_LLVM_PASS_H
+
+#include "mlir/Pass/Pass.h"
+#include <memory>
+
+namespace xilinx {
+namespace air {
+
+std::unique_ptr<mlir::Pass> createAIRTranslateToLLVMPass();
+
+} // namespace air
+} // namespace xilinx
+
+#endif // AIR_CONVERSION_AIR_TRANSLATE_TO_LLVM_PASS_H

mlir/include/air/Conversion/GPUPassDetail.h

@@ -23,6 +23,7 @@ namespace air {
 using namespace mlir;
 
 #define GEN_PASS_DECL
+#define GEN_PASS_DEF_AIRTRANSLATETOLLVM
 #define GEN_PASS_DEF_CONVERTAIRTOROCDL
 #define GEN_PASS_DEF_CONVERTGPUKERNELOUTLINE
 #include "air/Conversion/GPUPasses.h.inc"

mlir/include/air/Conversion/GPUPasses.td

@@ -21,6 +21,23 @@ def ConvertAIRToROCDL : Pass<"air-to-rocdl", "ModuleOp"> {
   let options = [];
 }
 
+def AIRTranslateToLLVM : Pass<"air-translate-to-llvm", "ModuleOp"> {
+  let summary = "Lower air.translate to memref.reinterpret_cast + LLVM-dialect address arithmetic";
+  let description = [{
+    Expands each `air.translate` op into the pointer-rebase computation:
+    `bases[to_rank] - bases[from_rank]`, converted from bytes to elements
+    of the source memref's element type, then applied as a new offset
+    via `memref.reinterpret_cast`. The expansion is pure arithmetic; it
+    works identically on host functions and inside `gpu.func`.
+  }];
+  let constructor = "xilinx::air::createAIRTranslateToLLVMPass()";
+  let dependentDialects = [
+    "mlir::arith::ArithDialect",
+    "mlir::memref::MemRefDialect",
+    "mlir::LLVM::LLVMDialect"
+  ];
+}
+
 def ConvertGPUKernelOutline : Pass<"air-gpu-outlining", "ModuleOp"> {
   let summary = "Outline GPU Kernel Func from GPU Launch";
   let constructor = "xilinx::air::createGPUKernelOutlinePass()";

mlir/include/air/Dialect/AIR/AIR.td

@@ -16,6 +16,14 @@ include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/TilingInterface.td"
 
+// Type predicate for !llvm.ptr. Inlined here (instead of including
+// "mlir/Dialect/LLVMIR/LLVMOpBase.td") to avoid pulling the LLVM dialect
+// into our TableGen scope — that would confuse `mlir-tblgen
+// -gen-dialect-doc` which expects exactly one dialect per .td file.
+def air_LLVMPtr : Type<CPred<"::llvm::isa<::mlir::LLVM::LLVMPointerType>($_self)">,
+                       "LLVM pointer",
+                       "::mlir::LLVM::LLVMPointerType">;
+
 class air_Op<string mnemonic, list<Trait> traits = []> :
     Op<air_Dialect, mnemonic, traits>;
 
@@ -923,6 +931,42 @@ def air_ExecuteTerminatorOp : air_Op<"execute_terminator", [HasParent<"ExecuteOp
       [{  attr-dict ($results^ `:` type($results))? }];
 }
 
+def air_TranslateOp : air_Op<"translate",
+                              [Pure, AllTypesMatch<["source", "result"]>]>,
+                       Arguments<(ins AnyMemRef:$source,
+                                      Index:$from_rank,
+                                      Index:$to_rank,
+                                      air_LLVMPtr:$heap_bases)>,
+                       Results<(outs AnyMemRef:$result)> {
+  let summary = "Re-express a symmetric-heap memref in another rank's address space";
+  let description = [{
+    Produces a memref of the same type as `$source` whose underlying
+    pointer references the corresponding allocation on `$to_rank`. The
+    `$source` memref is assumed to live on `$from_rank`'s symmetric heap.
+    The translation is the pointer rebase
+
+        peer_va = bases[to_rank] + (source_ptr - bases[from_rank])
+
+    where `$heap_bases` is the per-rank base table obtained from the
+    `mgpuGetHeapBases()` runtime hook (typically called once at host
+    scope and threaded through `gpu.launch_func` as a kernel argument).
+    No data is moved; this op produces a value-level "view" of peer
+    memory.
+
+    Folds to `$source` when `$from_rank` and `$to_rank` are statically
+    equal.
+
+    Both ranks must address the same collective allocation on the
+    symmetric heap (i.e. `$source` must trace back to a
+    `memref.alloc {air.symmetric}`). Using this op outside that contract
+    is undefined.
+  }];
+  let assemblyFormat =
+      [{ $source `,` $from_rank `,` $to_rank `,` $heap_bases
+         attr-dict `:` type($source) `,` type($heap_bases) }];
+  let hasFolder = 1;
+}
+
 // AIR custom op, as a handle for a user-provided AIE kernel
 
 def air_CustomOp : air_Op<"custom", [air_AsyncOpInterface,

mlir/include/air/Dialect/AIR/AIRDialect.h

@@ -9,6 +9,7 @@
 #ifndef MLIR_AIR_DIALECT_H
 #define MLIR_AIR_DIALECT_H
 
+#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"

mlir/lib/Conversion/AIRTranslateToLLVMPass.cpp

@@ -0,0 +1,179 @@
+//===- AIRTranslateToLLVMPass.cpp -------------------------------*- C++ -*-===//
+//
+// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+//===-----------------------------------------------------------------------===//
+//
+// Lower air.translate to memref-descriptor construction over a peer-rebased
+// pointer.
+//
+// For each `air.translate %src, %from, %to, %bases`:
+//   1. Extract the source memref's aligned pointer as !llvm.ptr.
+//   2. Compute the byte diff between the per-rank base pointers from the
+//      `$heap_bases` table:
+//         byte_diff = ptrtoint(bases[to]) - ptrtoint(bases[from])
+//   3. Apply the byte diff to the source aligned pointer (i8 GEP) to obtain
+//      the peer aligned pointer.
+//   4. Build a fresh LLVM memref descriptor (poison + insertvalue chain)
+//      whose allocated/aligned pointers both point at the peer address; the
+//      offset is 0, and sizes/strides are taken from the source memref's
+//      static type.
+//   5. unrealized_conversion_cast the descriptor back to the result memref
+//      type so downstream uses keep working through the standard
+//      memref-to-llvm pipeline.
+//
+// The lowering only uses arith + memref + llvm dialect ops — no runtime
+// calls. It is therefore valid both at host scope and inside `gpu.func`
+// (the kernel must already have been given the heap_bases pointer as a
+// kernel argument).
+//
+//===-----------------------------------------------------------------------===//
+
+#include "air/Conversion/AIRTranslateToLLVMPass.h"
+#include "air/Conversion/GPUPassDetail.h"
+#include "air/Dialect/AIR/AIRDialect.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+using namespace xilinx;
+
+namespace {
+
+// Build a fresh LLVM memref descriptor for `memrefTy` whose
+// allocated_ptr and aligned_ptr both reference `ptr`, offset is 0, and
+// sizes/strides come from the static type (row-major).
+//
+// Mirrors buildMemrefDescriptor in AIRSymmetricAllocToMgpuPass.
+static Value buildPeerDescriptor(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());
+
+  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) {
+    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 AIRTranslateToLLVMPass
+    : public xilinx::air::impl::AIRTranslateToLLVMBase<AIRTranslateToLLVMPass> {
+
+  AIRTranslateToLLVMPass() = default;
+  AIRTranslateToLLVMPass(const AIRTranslateToLLVMPass &) {}
+
+  void runOnOperation() override {
+    auto module = getOperation();
+    auto *ctx = module.getContext();
+    OpBuilder builder(ctx);
+    auto i64Ty = builder.getI64Type();
+    auto ptrTy = LLVM::LLVMPointerType::get(ctx);
+
+    SmallVector<air::TranslateOp> translates;
+    module.walk([&](air::TranslateOp op) { translates.push_back(op); });
+    if (translates.empty())
+      return;
+
+    for (air::TranslateOp op : translates) {
+      builder.setInsertionPoint(op);
+      Location loc = op.getLoc();
+
+      auto memrefTy = cast<MemRefType>(op.getSource().getType());
+      if (!memrefTy.hasStaticShape()) {
+        op.emitOpError("air.translate requires a static-shape source memref");
+        signalPassFailure();
+        return;
+      }
+
+      // Extract source aligned pointer as !llvm.ptr.
+      Value srcAlignedIdx = memref::ExtractAlignedPointerAsIndexOp::create(
+          builder, loc, op.getSource());
+      Value srcAlignedI64 = arith::IndexCastOp::create(builder, loc, i64Ty,
+                                                       srcAlignedIdx);
+      Value srcAlignedPtr =
+          LLVM::IntToPtrOp::create(builder, loc, ptrTy, srcAlignedI64);
+
+      // Load bases[from] and bases[to].
+      Value fromI64 = arith::IndexCastOp::create(builder, loc, i64Ty,
+                                                 op.getFromRank());
+      Value toI64 = arith::IndexCastOp::create(builder, loc, i64Ty,
+                                               op.getToRank());
+      Value fromBaseAddr = LLVM::GEPOp::create(
+          builder, loc, ptrTy, ptrTy, op.getHeapBases(), ValueRange{fromI64});
+      Value fromBase = LLVM::LoadOp::create(builder, loc, ptrTy, fromBaseAddr);
+      Value toBaseAddr = LLVM::GEPOp::create(builder, loc, ptrTy, ptrTy,
+                                             op.getHeapBases(),
+                                             ValueRange{toI64});
+      Value toBase = LLVM::LoadOp::create(builder, loc, ptrTy, toBaseAddr);
+
+      // byte_diff = ptrtoint(toBase) - ptrtoint(fromBase)
+      Value fromInt = LLVM::PtrToIntOp::create(builder, loc, i64Ty, fromBase);
+      Value toInt = LLVM::PtrToIntOp::create(builder, loc, i64Ty, toBase);
+      Value byteDiff = arith::SubIOp::create(builder, loc, toInt, fromInt);
+
+      // peer_aligned_ptr = srcAlignedPtr + byteDiff (as i8 GEP)
+      auto i8Ty = builder.getI8Type();
+      Value peerAlignedPtr = LLVM::GEPOp::create(
+          builder, loc, ptrTy, i8Ty, srcAlignedPtr, ValueRange{byteDiff});
+
+      // Build a fresh memref descriptor with the peer aligned pointer.
+      Value desc = buildPeerDescriptor(builder, loc, memrefTy, peerAlignedPtr);
+      Value newMemref =
+          UnrealizedConversionCastOp::create(builder, loc,
+                                             TypeRange{memrefTy},
+                                             ValueRange{desc})
+              .getResult(0);
+
+      op.getResult().replaceAllUsesWith(newMemref);
+      op.erase();
+    }
+  }
+};
+
+} // namespace
+
+namespace xilinx {
+namespace air {
+
+std::unique_ptr<mlir::Pass> createAIRTranslateToLLVMPass() {
+  return std::make_unique<AIRTranslateToLLVMPass>();
+}
+
+} // namespace air
+} // namespace xilinx

mlir/lib/Conversion/CMakeLists.txt

@@ -55,6 +55,7 @@ if(AIR_ENABLE_GPU)
   set(GPU_PASS_DEPENDS AIRGPUConversionPassIncGen)
   list(APPEND CONVERSION_SOURCES
     AIRToROCDLPass.cpp
+    AIRTranslateToLLVMPass.cpp
     GPUKernelOutlinePass.cpp
   )
   list(APPEND CONVERSION_LINK_LIBS

mlir/lib/Conversion/Passes.cpp

@@ -10,6 +10,7 @@
 
 #if AIR_ENABLE_GPU
 #include "air/Conversion/AIRToROCDLPass.h"
+#include "air/Conversion/AIRTranslateToLLVMPass.h"
 #include "air/Conversion/GPUKernelOutlinePass.h"
 #endif
 

mlir/lib/Dialect/AIR/IR/AIRDialect.cpp

@@ -3583,6 +3583,20 @@ ParseResult air::CustomOp::parse(OpAsmParser &parser, OperationState &result) {
   return success();
 }
 
+//
+// TranslateOp
+//
+
+OpFoldResult air::TranslateOp::fold(FoldAdaptor adaptor) {
+  if (getFromRank() == getToRank())
+    return getSource();
+  auto fromAttr = dyn_cast_if_present<IntegerAttr>(adaptor.getFromRank());
+  auto toAttr = dyn_cast_if_present<IntegerAttr>(adaptor.getToRank());
+  if (fromAttr && toAttr && fromAttr.getValue() == toAttr.getValue())
+    return getSource();
+  return {};
+}
+
 } // namespace xilinx
 
 #include "air/Dialect/AIR/AIROpInterfaces.cpp.inc"