Rewrite memory planning

include/PTO/IR/PTOAttrs.td

@@ -600,19 +600,6 @@ def PTO_EventAttr : PTO_Attr<"Event", "event"> {
   }];
 }
 
-//===----------------------------------------------------------------------===//
-// Plan Memory Mode
-//===----------------------------------------------------------------------===//
-
-def PTO_LOCALMEMPLAN : I32EnumAttrCase<"LOCAL_MEM_PLAN", 0>;
-def PTO_GLOBALWORKSPACEPLAN : I32EnumAttrCase<"GLOBAL_WORKSPACE_PLAN", 1>;
-
-def PTO_MemPlanModeEnum : PTO_I32Enum<
-  "MemPlanMode", "Mem Plan Mode", [
-    PTO_LOCALMEMPLAN,
-    PTO_GLOBALWORKSPACEPLAN
-]>;
-
 //===----------------------------------------------------------------------===//
 // Pad Mode for LoadOp
 //===----------------------------------------------------------------------===//

include/PTO/IR/PTOMultiBuffer.h

@@ -9,9 +9,10 @@
 //===- PTOMultiBuffer.h - Shared constants for multi-buffer ----*- C++ -*-===//
 //
 // Shared constants for the multi-buffer expression scheme:
-//   - `kPtoMultiBufferAttrName` is the memref-level attribute name written by
-//     PTOViewToMemref when lowering an `alloc_multi_tile` op. PlanMemory and
-//     downstream passes read it to reserve N physical slots.
+//   - `kPtoMultiBufferAttrName` is the legacy memref-level slot-count
+//     attribute.
+//   - `kPtoMultiBufferAddrsAttrName` is the internal address list written
+//     directly on tile-native `pto.alloc_multi_tile` by PlanMemory.
 //   - `kPtoMultiBufferMaxNum` is the upper bound on the slot count N. It is
 //     kept in lock-step with the InsertSync `MAX_MULTI_BUFFER_NUM`.
 //
@@ -29,6 +30,10 @@ namespace pto {
 inline constexpr llvm::StringLiteral kPtoMultiBufferAttrName =
     "pto.multi_buffer";
 
+/// Internal DenseI64ArrayAttr containing one byte address per physical slot.
+inline constexpr llvm::StringLiteral kPtoMultiBufferAddrsAttrName =
+    "pto.multi_buffer_addrs";
+
 /// Upper bound for N; must stay consistent with `MAX_MULTI_BUFFER_NUM` in
 /// insert-sync.
 inline constexpr unsigned kPtoMultiBufferMaxNum = 16;

include/PTO/IR/PTOOps.td

@@ -343,19 +343,18 @@ def AllocTileOp : PTO_Op<"alloc_tile", [AttrSizedOperandSegments]> {
 //===----------------------------------------------------------------------===//
 //
 // `pto.alloc_multi_tile` declares an N-slot logical tile buffer. The actual
-// physical allocation of N slots is performed by PTOPlanMemory once the op
-// has been lowered to `memref.alloc {pto.multi_buffer = N : i32}` by
-// PTOViewToMemref.
+// physical allocation of N slots is performed directly by PTOPlanMemory.
+// The op remains tile-native until PTOResolveBufferSelect materializes an
+// addressed `pto.alloc_tile` for each selected slot.
 //
 // `pto.multi_tile_get` is the *only* way to consume a `multi_tile_buf`: it
 // returns a regular `tile_buf` view onto the chosen slot. The slot index
 // may be a constant or any SSA `index` value; PTOAS does not synthesize
 // `iv mod N` automatically -- the user expression IS the slot selector.
 //
-// `pto.slot_marker` is an internal op materialized by PTOViewToMemref to
-// thread the slot SSA from `multi_tile_get` through the memref layer down
-// to PlanMemory / sync analysis. It is not intended for direct frontend
-// use.
+// `pto.slot_marker` remains as an internal compatibility op for legacy
+// memref-based multi-buffer IR. PTOViewToMemref does not create it for
+// tile-native `pto.multi_tile_get`.
 
 def AllocMultiTileOp : PTO_Op<"alloc_multi_tile", [AttrSizedOperandSegments]> {
   let summary = "Allocate an N-slot multi-buffer tile";
@@ -415,9 +414,9 @@ def MultiTileGetOp : PTO_Op<"multi_tile_get", [
     sync analysis treats this use as a slot-indexed dynamic access; ptoas
     does NOT rewrite the user expression into `iv mod N`.
 
-    This op is metadata-only (no data movement); the lowering merely
-    annotates the underlying memref view with the slot index for PlanMemory
-    / sync / EnableBufferSelect to consume.
+    This op is metadata-only (no data movement). PlanMemory preserves it,
+    sync analysis consumes its slot index directly, and
+    PTOResolveBufferSelect replaces it with an addressed `pto.alloc_tile`.
   }];
 
   let arguments = (ins
@@ -560,7 +559,7 @@ def SetValidShapeOp : PTO_Op<"set_validshape", [
   }];
 
   let arguments = (ins
-    TileBufOrMemRef:$source,
+    TileBufType:$source,
     Index:$valid_row,
     Index:$valid_col
   );
@@ -584,7 +583,7 @@ def GetValidShapeOp : PTO_Op<"get_validshape", [
   }];
 
   let arguments = (ins
-    TileBufOrMemRef:$source
+    TileBufType:$source
   );
 
   let results = (outs
@@ -1556,11 +1555,10 @@ def SlotMarkerOp : PTO_Op<"slot_marker", [
   ]> {
   let summary = "Tag a memref view as referring to one slot of a multi_tile_buf";
   let description = [{
-    Internal op materialized by `PTOViewToMemref` while lowering
-    `pto.multi_tile_get`. It carries the slot SSA index through the memref
-    layer so that PlanMemory, sync analysis (InsertSync / GraphSyncSolver),
-    and the buffer-select lowering pass can identify which physical slot
-    this memref reference touches.
+    Internal compatibility op for legacy memref-based multi-buffer IR. It
+    carries a slot SSA index so sync analysis (InsertSync / GraphSyncSolver)
+    and the buffer-select lowering pass can identify which physical slot a
+    memref reference touches.
 
     The op is metadata-only (no data movement, no extra address arithmetic);
     its result memref aliases the source memref byte-for-byte. Frontends do
@@ -2634,12 +2632,10 @@ def LocalArraySetOp : PTO_Op<"local_array_set"> {
 }
 
 def DeclareTileMemRefOp : PTO_Op<"declare_tile_memref"> {
-  let summary = "Internal memref placeholder for a tile whose address is assigned later";
+  let summary = "Legacy memref placeholder for a runtime-bound tile";
   let description = [{
-    Internal lowering op used by PTOViewToMemref. This op does not allocate
-    storage; it only provides a memref-typed SSA handle so later passes can
-    attach tile metadata through pto.bind_tile before the address is filled by
-    pipe operations such as pto.tpop.
+    Compatibility op for legacy memref-based IR. New PTOViewToMemref pipelines
+    preserve `pto.declare_tile` directly instead of producing this op.
   }];
 
   let results = (outs AnyMemRef:$result);

include/PTO/IR/VPTOOps.td

@@ -220,11 +220,10 @@ def TensorViewAddrOp : PTO_Op<"tensor_view_addr", [Pure]> {
     typed PTO pointer, depending on the requested destination type.
 
     In authoring-form IR this op preserves the descriptor-style surface;
-    during view-to-memref lowering it collapses to the underlying memref value
-    or to a memref-derived pointer.
+    tile-native inputs remain unchanged through memory planning and sync.
 
-    This op may also accept a memref operand after earlier view lowering, in
-    which case it behaves as an identity marker and is removed by lowering.
+    This op may also accept a legacy memref operand, in which case it behaves
+    as an identity marker and is removed by compatibility lowering.
   }];
 
   let arguments = (ins AnyTypeOf<[TensorViewType, PartitionTensorViewType, AnyMemRef],
@@ -247,9 +246,9 @@ def TileBufAddrOp : PTO_Op<"tile_buf_addr", [Pure]> {
     shape and address space.
 
     This op is emitted by TileLang DSL templates and resolved by the
-    FoldTileBufIntrinsics pass after inlining. Hand-written `.pto` may also
-    use it directly on the memref result of `pto.bind_tile` / lowered
-    `pto.alloc_tile`.
+    FoldTileBufIntrinsics pass after inlining. Tile-native uses remain in PTO
+    IR through memory planning and sync. Hand-written legacy `.pto` may also
+    use it directly on the memref result of `pto.bind_tile`.
   }];
 
   let arguments = (ins AnyTypeOf<[TileBufType, AnyMemRef],

include/PTO/Transforms/InsertSync/PTOIRTranslator.h

@@ -70,6 +70,8 @@ class PTOIRTranslator {
   // --- 内存/Alias 分析 ---
   void UpdateKernelArgMemInfo();
   LogicalResult UpdateAllocTileOpMemInfo(pto::AllocTileOp op);
+  LogicalResult UpdateAllocMultiTileOpMemInfo(pto::AllocMultiTileOp op);
+  LogicalResult UpdateDeclareTileOpMemInfo(pto::DeclareTileOp op);
   LogicalResult UpdateDeclareGlobalOpMemInfo(pto::DeclareGlobalOp op);
   LogicalResult UpdateDeclareTileMemRefOpMemInfo(pto::DeclareTileMemRefOp op);
   LogicalResult UpdatePointerCastOpMemInfo(pto::PointerCastOp op);
@@ -81,6 +83,9 @@ class PTOIRTranslator {
   void UpdateMemrefSubViewAliasBufferInfo(memref::SubViewOp op);
   void UpdateTileSubViewAliasBufferInfo(pto::SubViewOp op);
   void UpdateSlotMarkerAliasBufferInfo(pto::SlotMarkerOp op);
+  void UpdateMultiTileGetAliasBufferInfo(pto::MultiTileGetOp op);
+  void UpdateSlotSelectedAliasBufferInfo(Value result, Value source,
+                                         Value slot);
 
   // --- 控制流处理 (SCF) ---
   void UpdateForOpInfo(scf::ForOp forOp);

include/PTO/Transforms/Passes.h

@@ -67,10 +67,10 @@ std::unique_ptr<Pass> createConvertToPTOOpPass();
 /// PTO Ops.
 std::unique_ptr<Pass> createInferPTOMemScopePass();
 
-/// Create a pass to plan memory.
 std::unique_ptr<Pass>
-createPlanMemoryPass(const PlanMemoryOptions &planMemoryOption = {});
-
+createPlanMemoryPass(const PlanMemoryOptions &options = {});
+std::unique_ptr<Pass>
+createPlanMemoryModernPass(const PlanMemoryOptions &options);
 std::unique_ptr<Pass> createPTORemoveRedundantBarrierPass();
 std::unique_ptr<Pass> createPTOViewToMemrefPass();
 std::unique_ptr<Pass> createPTOValidateIntToPtrUsesPass();

include/PTO/Transforms/Passes.td

@@ -193,37 +193,26 @@ def PTORematerializeFixpipeVectorQuant
 }
 
 def PlanMemory : Pass<"pto-plan-memory", "ModuleOp"> {
-  let summary = "Plan memory for PTO Ops";
+ let summary = "Plan memory for PTO Ops";
   let constructor = "mlir::pto::createPlanMemoryPass()";
-  let dependentDialects = ["pto::PTODialect", ];
+
   let options = [
-    Option<"memMode", "mem-plan-mode", "pto::MemPlanMode",
-           "pto::MemPlanMode::LOCAL_MEM_PLAN",
-           "plan mem mode (default is LOCAL_MEM_PLAN)",
-           [{::llvm::cl::values(
-               clEnumValN(pto::MemPlanMode::LOCAL_MEM_PLAN, "local-mem-plan",
-                          "plan mem mode is for memref.alloc"),
-               clEnumValN(
-                   pto::MemPlanMode::GLOBAL_WORKSPACE_PLAN,
-                   "global-work-space-plan",
-                   "plan mem mode is for memref_ext.alloc_workspace"))}]>,
-    Option<"enableGlobalReuse", "enable-global-workspace-reuse", "bool",
-           /*default=*/"false",
-           "Enable global workspace reuse ,default : false">,
-    Option<"enablePrintMemoryAllocatedSize", "enable-print-memory-allocated-size", "bool",
-           /*default=*/"false",
-           "Enable print memory allocated size, default : false">,
-    Option<"restrictInplaceAsISA", "restrict-inplace-as-isa", "bool",
-           /*default=*/"false",
-           "restrict memory inplace as isa, default : false">,
-    Option<"orderBySize", "order-by-size", "bool",
-           /*default=*/"false",
-           "Process buffers largest-first (first-fit-decreasing order) during "
-           "local memory planning instead of the default DMA-first order. "
-           "Decreasing-size order packs heterogeneous-size buffers tighter "
-           "(matches XLA/TVM/SOMAS). default : false">,
+    Option<"memMode", "mem-mode", "std::string", "\"local\"",
+           "Planning mode. Supported value today: local">,
+    Option<"orderBySize", "order-by-size", "bool", "false",
+           "Plan larger buffers first inside one AddressSpace before applying "
+           "SPEC_LEVEL_0 first-fit reuse">
+  ];
+
+  let dependentDialects = [
+    "mlir::pto::PTODialect",
+    "mlir::memref::MemRefDialect",
+    "mlir::arith::ArithDialect",
+    "mlir::func::FuncDialect",
+    "mlir::scf::SCFDialect"
   ];
 }
+
 def PTOLoweringSyncToPipe : Pass<"pto-lowering-sync-to-pipe", "func::FuncOp"> {
   let summary = "Lower high-level sync ops to low-level pipe ops";
   let description = [{
@@ -752,25 +741,20 @@ def PTOMaterializeTileHandles : Pass<"pto-materialize-tile-handles", "ModuleOp">
 }
 
 def PTOResolveBufferSelect : Pass<"pto-resolve-buffer-select", "ModuleOp"> {
-  let summary = "Lower pto.slot_marker to single-slot pointer_cast (multi-buffer)";
+  let summary = "Resolve multi-buffer slot selections to addressed tile handles";
   let description = [{
-    Consumes `pto.slot_marker %src[%k]` ops written by PTOViewToMemref to
-    thread multi-buffer slot selection through the memref layer.
-
-    The op is replaced by a fresh single-address `pto.pointer_cast` that
-    refers to the chosen physical slot:
+    Primarily consumes tile-native `pto.multi_tile_get %src[%k]` operations.
+    The source `pto.alloc_multi_tile` carries either planner-assigned physical
+    slot addresses or a level3 base address. Each selection is replaced by a
+    `pto.alloc_tile` carrying the selected address; constant slots select one
+    address directly, while dynamic slots use an N-way `arith.select` chain.
 
-    - constant slot `k`: a single `pto.pointer_cast(addrK)` is emitted at
-      the use site, using slot k from the underlying multi-address
-      pointer_cast created by PTOPlanMemory / AllocToPointerCast.
-    - dynamic slot `%k`: per-slot single-address `pto.pointer_cast`s are
-      created and an N-way `arith.select` chain picks one according to the
-      user-supplied SSA. This pass does NOT rewrite `%k` into `iv mod N`;
-      the frontend expression is the slot selector.
+    The pass also retains compatibility with legacy memref-based
+    `pto.slot_marker` IR, where the selection is materialized as a
+    single-address `pto.pointer_cast`.
 
-    The original multi-address `pto.pointer_cast` is kept in IR as the
-    "alloc anchor" so future sync passes can recognize the multi-buffer
-    geometry.
+    This pass does not rewrite `%k` into `iv mod N`; the frontend expression
+    remains the slot selector.
   }];
 
   let constructor = "mlir::pto::createPTOResolveBufferSelectPass()";