[IR] Verify explicit local tile address alignment

include/PTO/IR/PTOOps.td

@@ -1425,6 +1425,10 @@ def PointerCastOp : PTO_Op<"pointer_cast", [AttrSizedOperandSegments, Pure]> {
         $_builder.getDenseI32ArrayAttr({addrsSize, vRowSize, vColSize}));
     }]>
   ];
+
+  let extraClassDeclaration = [{
+    ::mlir::LogicalResult verify();
+  }];
 }
 
 // =============================================================================

lib/PTO/IR/PTO.cpp

@@ -2811,12 +2811,87 @@ void mlir::pto::annotatePTOEntryFunctions(ModuleOp module) {
   return success();
 }
 
+static std::optional<uint64_t>
+getLocalAddressAlignmentBytes(Attribute memorySpace) {
+  auto addrSpace = dyn_cast_or_null<AddressSpaceAttr>(memorySpace);
+  if (!addrSpace)
+    return std::nullopt;
+
+  switch (addrSpace.getAddressSpace()) {
+  case AddressSpace::VEC:
+  case AddressSpace::MAT:
+  case AddressSpace::BIAS:
+  case AddressSpace::SCALING:
+    return 32;
+  case AddressSpace::LEFT:
+  case AddressSpace::RIGHT:
+  case AddressSpace::ACC:
+    return 512;
+  case AddressSpace::GM:
+  case AddressSpace::Zero:
+    return std::nullopt;
+  }
+  return std::nullopt;
+}
+
+static std::optional<int64_t> getConstantIntValue(Value value) {
+  if (!value)
+    return std::nullopt;
+
+  auto constOp = value.getDefiningOp<arith::ConstantOp>();
+  if (!constOp)
+    return std::nullopt;
+
+  auto intAttr = dyn_cast<IntegerAttr>(constOp.getValue());
+  if (!intAttr)
+    return std::nullopt;
+
+  return intAttr.getValue().getSExtValue();
+}
+
+static LogicalResult verifyConstantLocalAddress(Operation *op, Value addr,
+                                                Attribute memorySpace,
+                                                int addrIndex = -1) {
+  std::optional<uint64_t> alignment =
+      getLocalAddressAlignmentBytes(memorySpace);
+  if (!alignment || *alignment == 0)
+    return success();
+
+  std::optional<int64_t> constantAddr = getConstantIntValue(addr);
+  if (!constantAddr)
+    return success();
+
+  auto emitAddrError = [&]() {
+    InFlightDiagnostic diag = op->emitOpError();
+    if (addrIndex >= 0)
+      diag << "addr[" << addrIndex << "]";
+    else
+      diag << "addr";
+    return diag;
+  };
+
+  if (*constantAddr < 0)
+    return emitAddrError() << " must be non-negative, got " << *constantAddr;
+
+  uint64_t unsignedAddr = static_cast<uint64_t>(*constantAddr);
+  if ((unsignedAddr % *alignment) != 0)
+    return emitAddrError()
+           << " must be aligned to " << *alignment
+           << " bytes for local tile memory space, got " << unsignedAddr;
+
+  return success();
+}
+
 LogicalResult AllocTileOp::verify() {
   auto ty = getResult().getType(); // TileBufType
 
   if (failed(verifyTileBufLayoutConstraints(*this, ty, "result")))
     return failure();
 
+  if (failed(verifyConstantLocalAddress(getOperation(), getAddr(),
+                                        ty.getMemorySpace())))
+    return failure();
+
   // op 上有没有传 operands
   bool hasVR = getValidRow() != nullptr;
   bool hasVC = getValidCol() != nullptr;
@@ -2848,6 +2923,21 @@ LogicalResult AllocTileOp::verify() {
   return success();
 }
 
+LogicalResult PointerCastOp::verify() {
+  auto memRefTy = dyn_cast<BaseMemRefType>(getResult().getType());
+  if (!memRefTy)
+    return emitOpError("result must be a memref type");
+
+  for (auto [idx, addr] : llvm::enumerate(getAddrs())) {
+    if (failed(verifyConstantLocalAddress(getOperation(), addr,
+                                          memRefTy.getMemorySpace(),
+                                          static_cast<int>(idx))))
+      return failure();
+  }
+
+  return success();
+}
+
 LogicalResult MaterializeTileOp::verify() {
   auto sourceTy = cast<MemRefType>(getSource().getType());
   auto resultTy = cast<TileBufType>(getResult().getType());

test/basic/alloc_tile_addr_alignment_invalid.mlir

@@ -0,0 +1,11 @@
+// RUN: not ptoas --pto-level=level3 %s 2>&1 | FileCheck %s
+
+module {
+  func.func @unaligned_alloc_tile_addr() {
+    %bad = arith.constant 70660 : i64
+    // CHECK: pto.alloc_tile
+    // CHECK: addr must be aligned to 32 bytes for local tile memory space, got 70660
+    %tile = pto.alloc_tile addr = %bad : !pto.tile_buf<loc=vec, dtype=f16, rows=32, cols=64, v_row=32, v_col=64, blayout=row_major, slayout=none_box, fractal=512, pad=0>
+    return
+  }
+}

test/basic/alloc_tile_addr_alignment_valid.mlir

@@ -0,0 +1,9 @@
+// RUN: ptoas --pto-level=level3 %s >/dev/null
+
+module {
+  func.func @aligned_alloc_tile_addr() {
+    %ok = arith.constant 71232 : i64
+    %tile = pto.alloc_tile addr = %ok : !pto.tile_buf<loc=vec, dtype=f16, rows=32, cols=64, v_row=32, v_col=64, blayout=row_major, slayout=none_box, fractal=512, pad=0>
+    return
+  }
+}

test/basic/pointer_cast_addr_alignment_gm_ignored.mlir

@@ -0,0 +1,9 @@
+// RUN: ptoas --pto-level=level3 %s >/dev/null
+
+module {
+  func.func @gm_pointer_cast_addr_is_ignored() {
+    %neg = arith.constant -1 : i64
+    %buf = pto.pointer_cast(%neg) : memref<32x64xf16, #pto.address_space<gm>>
+    return
+  }
+}

test/basic/pointer_cast_addr_alignment_invalid.mlir

@@ -0,0 +1,11 @@
+// RUN: not ptoas --pto-level=level3 %s 2>&1 | FileCheck %s
+
+module {
+  func.func @unaligned_pointer_cast_addr() {
+    %bad = arith.constant 70660 : i64
+    // CHECK: pto.pointer_cast
+    // CHECK: addr[0] must be aligned to 32 bytes for local tile memory space, got 70660
+    %buf = pto.pointer_cast(%bad) : memref<32x64xf16, #pto.address_space<vec>>
+    return
+  }
+}