feat(vpto): prepare PTODSL pipe expansion

include/PTO/IR/PTOOps.td

@@ -40,6 +40,14 @@ def TensorOrMemref :
 def PTODpsType :
   AnyTypeOf<[AnyRankedTensor, AnyMemRef, PartitionTensorViewType, TileBufType]>;
 
+// TASSIGN also rebases a declared GlobalTensor descriptor. Keep this separate
+// from PTODpsType: ordinary DPS operations must not accept a raw TensorViewType
+// in place of their logical partition view.
+def PTOTAssignType :
+  AnyTypeOf<[AnyRankedTensor, AnyMemRef, TensorViewType,
+             PartitionTensorViewType, TileBufType],
+            "TensorView, PartitionTensorView, TileBuf, Tensor, or MemRef">;
+
 def PTOPipeEntryType :
   AnyTypeOf<[AnyRankedTensor, AnyMemRef, TensorViewType, TileBufType],
             "TensorView, TileBuf, Tensor, or MemRef">;
@@ -2437,6 +2445,7 @@ def InitializeL2LPipeOp : PTO_Op<"initialize_l2l_pipe", [
 }
 
 def TPushOp : PTO_TOp<"tpush", [
+    AttrSizedOperandSegments,
     OpPipeInterface,
     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
 ]> {
@@ -2446,6 +2455,7 @@ def TPushOp : PTO_TOp<"tpush", [
       PTOPipeEntryType:$tile,
       PipeType:$pipe_handle,
       Optional<I64>:$aiv_subblockid,
+      Optional<StructType>:$pipe_state,
       I8Attr:$split
   );
 
@@ -2456,6 +2466,7 @@ def TPushOp : PTO_TOp<"tpush", [
     `(` $tile `,` $pipe_handle `:` qualified(type($tile)) `,` qualified(type($pipe_handle)) `)`
     `{` `split` `=` $split `}`
     (`aiv_subblockid` `(` $aiv_subblockid^ `)`)?
+    (`pipe_state` `(` $pipe_state^ `:` qualified(type($pipe_state)) `)`)?
     attr-dict
   }];
 
@@ -2519,6 +2530,7 @@ def TAllocOp : PTO_TOp<"talloc", [
   let arguments = (ins
       TensorViewType:$entry,
       PipeType:$pipe_handle,
+      Optional<StructType>:$pipe_state,
       I8Attr:$split
   );
 
@@ -2527,17 +2539,17 @@ def TAllocOp : PTO_TOp<"talloc", [
 
   let assemblyFormat = [{
     `(` $entry `,` $pipe_handle `:` qualified(type($entry)) `,` qualified(type($pipe_handle)) `)`
-    `{` `split` `=` $split `}` attr-dict
+    `{` `split` `=` $split `}`
+    (`pipe_state` `(` $pipe_state^ `:` qualified(type($pipe_state)) `)`)?
+    attr-dict
   }];
 
   let extraClassDeclaration = [{
     ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_MTE3; }
   }];
 }
 
-def TAssignOp : PTO_Op<"tassign", [
-    Pure
-  ]> {
+def TAssignOp : PTO_Op<"tassign"> {
   let summary = "Rebind a declared tile handle to a runtime address";
   let description = [{
     Rebinds an existing tile handle to a new address. This op is intended for
@@ -2546,10 +2558,10 @@ def TAssignOp : PTO_Op<"tassign", [
   }];
 
   let arguments = (ins
-    PTODpsType:$tile,
+    PTOTAssignType:$tile,
     I64:$addr
   );
-  let results = (outs PTODpsType:$result);
+  let results = (outs PTOTAssignType:$result);
 
   let assemblyFormat = [{
     $tile `,` $addr attr-dict `:` qualified(type($tile)) `->` qualified(type($result))
@@ -2790,6 +2802,7 @@ def DeclareTileMemRefOp : PTO_Op<"declare_tile_memref"> {
 
 def TPopOp : PTO_TOp<"tpop", [
     PTO_DpsInitOpInterface,
+    AttrSizedOperandSegments,
     OpPipeInterface,
     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
 ]> {
@@ -2799,6 +2812,7 @@ def TPopOp : PTO_TOp<"tpop", [
       PTOPipeEntryType:$tile,
       PipeType:$pipe_handle,
       Optional<I64>:$aiv_subblockid,
+      Optional<StructType>:$pipe_state,
       I8Attr:$split
   );
 
@@ -2809,6 +2823,7 @@ def TPopOp : PTO_TOp<"tpop", [
     `(` $tile `,` $pipe_handle `:` qualified(type($tile)) `,` qualified(type($pipe_handle)) `)`
     `{` `split` `=` $split `}`
     (`aiv_subblockid` `(` $aiv_subblockid^ `)`)?
+    (`pipe_state` `(` $pipe_state^ `:` qualified(type($pipe_state)) `)`)?
     attr-dict
   }];
 
@@ -2864,6 +2879,7 @@ def TPopOp : PTO_TOp<"tpop", [
 }
 
 def TFreeOp : PTO_TOp<"tfree", [
+    AttrSizedOperandSegments,
     OpPipeInterface,
     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
 ]> {
@@ -2872,6 +2888,7 @@ def TFreeOp : PTO_TOp<"tfree", [
   let arguments = (ins
       Optional<PTOPipeEntryType>:$entry,
       PipeType:$pipe_handle,
+      Optional<StructType>:$pipe_state,
       I8Attr:$split
   );
 
@@ -2884,6 +2901,35 @@ def TFreeOp : PTO_TOp<"tfree", [
   }];
 }
 
+def TDrainOp : PTO_TOp<"tdrain", [
+    OpPipeInterface,
+    DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
+]> {
+  let summary = "Complete producer-side cleanup for a stateful unified pipe";
+  let description = [{
+    Internal operation materialized only by the VPTO PipeState path. It models
+    the terminal producer-side synchronization performed by the EmitC `TPipe`
+    object during destruction.
+  }];
+
+  let arguments = (ins
+      PipeType:$pipe_handle,
+      StructType:$pipe_state,
+      I8Attr:$split
+  );
+  let results = (outs);
+  let hasVerifier = 1;
+
+  let assemblyFormat = [{
+    `(` $pipe_handle `,` $pipe_state `:` qualified(type($pipe_handle)) `,` qualified(type($pipe_state)) `)`
+    `{` `split` `=` $split `}` attr-dict
+  }];
+
+  let extraClassDeclaration = [{
+    ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_MTE3; }
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // Synchronization Ops
 //===----------------------------------------------------------------------===//

include/PTO/IR/PTOTypeUtils.h

@@ -36,6 +36,10 @@ inline LogicalResult getPTOMemRefStridesAndOffset(
   return detail::getPTOMemRefStridesAndOffsetImpl(memTy, strides, offset, 0);
 }
 
+/// Returns whether `value` is a pto.declare_global result or a result of one
+/// or more pto.tassign operations rooted at that declaration.
+bool isDeclaredGlobalOrTAssignResult(Value value);
+
 bool isPTOFloat8Type(Type t);
 bool isPTOFloat8E4M3LikeType(Type t);
 bool isPTOFloat8E5M2LikeType(Type t);

include/PTO/Transforms/Passes.h

@@ -37,6 +37,7 @@ std::unique_ptr<Pass> createLoweringSyncToPipePass();
 std::unique_ptr<Pass> createPTOAssignDefaultFrontendPipeIdPass();
 std::unique_ptr<Pass> createPTOLowerFrontendPipeOpsPass();
 std::unique_ptr<Pass> createPTOInferValidatePipeInitPass();
+std::unique_ptr<Pass> createPTOMaterializePipeStatePass();
 std::unique_ptr<Pass> createPTOResolveReservedBuffersPass();
 std::unique_ptr<Pass> createPTOWrapFunctionsInSectionsPass();
 std::unique_ptr<Pass> createPTONormalizeUncoveredTileSectionsPass();

include/PTO/Transforms/Passes.td

@@ -484,6 +484,25 @@ def PTOInferValidatePipeInit : Pass<"pto-infer-validate-pipe-init", "ModuleOp">
   ];
 }
 
+def PTOMaterializePipeState : Pass<"pto-materialize-pipe-state", "func::FuncOp"> {
+  let summary = "Materialize explicit PipeState operands for unified pipe ops";
+  let description = [{
+    Creates one `!pto.struct<i32, i32>` state object per internal pipe handle,
+    initializes its producer and consumer counters to zero, attaches it to
+    unified `talloc` / `tpush` / `tpop` / `tfree` operations, and inserts a
+    terminal `tdrain` for every pipe with producer pushes. The terminal drain
+    follows the resolved pipe-level `nosplit` configuration (`split = 0` when
+    true, otherwise `split = 1`), so producer operations on a split-capable
+    pipe may use different split axes.
+  }];
+  let constructor = "mlir::pto::createPTOMaterializePipeStatePass()";
+  let dependentDialects = [
+    "mlir::arith::ArithDialect",
+    "mlir::pto::PTODialect",
+    "mlir::func::FuncDialect"
+  ];
+}
+
 def PTOResolveReservedBuffers : Pass<"pto-resolve-reserved-buffers", "ModuleOp"> {
   let summary = "Resolve reserved local buffer addresses and peer pipe flag bases";
   let description = [{
@@ -533,7 +552,13 @@ def InsertTemplateAttributes
            "PYTHONPATH root for PTODSL">,
     Option<"daemonHelperModule", "daemon-helper-module", "std::string",
            /*default=*/"\"ptodsl.tilelib.serving.helper\"",
-           "Python module used for daemon metadata RPC calls">
+           "Python module used for daemon metadata RPC calls">,
+    Option<"pipeOnly", "pipe-only", "bool",
+           /*default=*/"false",
+           "Attach candidates only to unified internal pipe operations">,
+    Option<"skipPipeOps", "skip-pipe-ops", "bool",
+           /*default=*/"false",
+           "Exclude unified internal pipe operations from this metadata pass">
   ];
 }