[ttl] Pipe verification; implement recv post -> send post -> send wait -> recv wait protocol (#608)

README.md

@@ -140,6 +140,7 @@ docker run -d --name $USER-ird \
   -v /dev/hugepages-1G:/dev/hugepages-1G \
   -v $HOME:$HOME \
   -v $SSH_AUTH_SOCK:/ssh-agent -e SSH_AUTH_SOCK=/ssh-agent \
+  -v $HOME/.ssh/known_hosts:/root/.ssh/known_hosts:ro \
   ghcr.io/tenstorrent/tt-lang/tt-lang-ird-ubuntu-22-04:latest \
   sleep infinity
 ```

include/ttlang/Dialect/TTL/IR/TTLOps.td

@@ -599,6 +599,42 @@ def TTL_IfDstOp : TTL_Op<"if_dst",
   }];
 }
 
+def TTL_PipeRecvPostOp : TTL_Op<"pipe_recv_post",
+    [MemoryEffects<[MemRead, MemWrite]>]> {
+  let summary = "Post a pipe receive destination";
+  let description = [{
+    Publishes the destination dataflow buffer slot address for a pipe receive
+    and returns a transfer handle. The operation does not wait for the payload.
+    A matching `ttl.pipe_recv_wait` completes the receive before the destination
+    slot is consumed.
+  }];
+  let arguments = (ins
+    TTL_Pipe:$pipe,
+    AnyType:$dst
+  );
+  let results = (outs TTL_TransferHandle:$xf);
+  let assemblyFormat = "$pipe `,` $dst attr-dict `:` functional-type(operands, results)";
+  let hasVerifier = 1;
+}
+
+def TTL_PipeRecvWaitOp : TTL_Op<"pipe_recv_wait",
+    [MemoryEffects<[MemRead, MemWrite]>]> {
+  let summary = "Wait for a posted pipe receive";
+  let description = [{
+    Blocks until the sender has completed the pipe transfer associated with a
+    posted receive. The destination dataflow buffer slot is safe to consume
+    after this operation completes. The `dst` operand identifies the posted
+    destination slot whose receive is being completed.
+  }];
+  let arguments = (ins
+    TTL_TransferHandle:$xf,
+    TTL_Pipe:$pipe,
+    AnyType:$dst
+  );
+  let assemblyFormat = "$xf `,` $pipe `,` $dst attr-dict `:` type($xf) `,` type($pipe) `,` type($dst)";
+  let hasVerifier = 1;
+}
+
 def TTL_PipeNetScopeOp : TTL_Op<"pipenet_scope",
     [SingleBlockImplicitTerminator<"YieldOp">,
      DeclareOpInterfaceMethods<RegionBranchOpInterface>]> {

include/ttlang/Dialect/TTL/IR/TTLOpsUtils.h

@@ -16,8 +16,10 @@
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
 
 #include <cstdint>
 #include <optional>
@@ -51,15 +53,82 @@ inline mlir::Value traceUnrealizedCasts(mlir::Value value) {
   return value;
 }
 
+/// Trace a transfer handle through tensor containers and loop-carried values
+/// to the first value accepted by `match`.
+template <typename ResultT, typename MatchFn>
+inline ResultT
+traceTransferHandleSource(mlir::Value value, MatchFn match,
+                          llvm::SmallPtrSetImpl<mlir::Value> &seen) {
+  value = traceUnrealizedCasts(value);
+  if (!seen.insert(value).second) {
+    return ResultT();
+  }
+
+  if (ResultT result = match(value)) {
+    return result;
+  }
+  if (auto extractOp = value.getDefiningOp<mlir::tensor::ExtractOp>()) {
+    return traceTransferHandleSource<ResultT>(extractOp.getTensor(), match,
+                                              seen);
+  }
+  if (auto insertOp = value.getDefiningOp<mlir::tensor::InsertOp>()) {
+    return traceTransferHandleSource<ResultT>(insertOp.getScalar(), match,
+                                              seen);
+  }
+  if (auto result = mlir::dyn_cast<mlir::OpResult>(value)) {
+    if (auto loop =
+            mlir::dyn_cast<mlir::LoopLikeOpInterface>(result.getOwner())) {
+      auto yieldedOpt = loop.getYieldedValuesMutable();
+      auto resultsOpt = loop.getLoopResults();
+      if (yieldedOpt && resultsOpt) {
+        auto yielded = *yieldedOpt;
+        auto results = *resultsOpt;
+        for (unsigned idx = 0; idx < results.size(); ++idx) {
+          if (results[idx] == result) {
+            return traceTransferHandleSource<ResultT>(yielded[idx].get(), match,
+                                                      seen);
+          }
+        }
+      }
+    }
+  }
+  if (auto blockArg = mlir::dyn_cast<mlir::BlockArgument>(value)) {
+    mlir::Operation *parent = blockArg.getOwner()->getParentOp();
+    if (auto loop = mlir::dyn_cast_or_null<mlir::LoopLikeOpInterface>(parent)) {
+      auto iterArgs = loop.getRegionIterArgs();
+      auto inits = loop.getInitsMutable();
+      for (unsigned idx = 0; idx < iterArgs.size(); ++idx) {
+        if (iterArgs[idx] == blockArg) {
+          return traceTransferHandleSource<ResultT>(inits[idx].get(), match,
+                                                    seen);
+        }
+      }
+    }
+  }
+
+  return ResultT();
+}
+
 /// Walk through `tensor.extract_slice` ops and return the underlying
 /// `ttl.cb_reserve` op, or null if the chain doesn't end at one.
 inline mlir::tt::ttl::CBReserveOp findCBReserveForView(mlir::Value view) {
+  view = traceUnrealizedCasts(view);
   while (auto slice = view.getDefiningOp<mlir::tensor::ExtractSliceOp>()) {
     view = slice.getSource();
+    view = traceUnrealizedCasts(view);
   }
   return view.getDefiningOp<mlir::tt::ttl::CBReserveOp>();
 }
 
+/// Return the user reserve that produced a pipe receive destination block.
+inline mlir::tt::ttl::CBReserveOp findCBReserveForPipeReceive(mlir::Value dst) {
+  dst = traceUnrealizedCasts(dst);
+  if (auto attach = dst.getDefiningOp<mlir::tt::ttl::AttachCBOp>()) {
+    return findCBReserveForView(attach.getTensor());
+  }
+  return findCBReserveForView(dst);
+}
+
 /// Resolve the CB index attached to `cb`, accepting either the pre-conversion
 /// BindCBOp or the post-conversion GetCompileArgValOp.
 inline std::optional<int64_t> getCBIndex(mlir::Value cb) {

include/ttlang/Dialect/TTL/Passes.td

@@ -520,6 +520,10 @@ def TTLVerifyPipeNetGuards
     Each diagnostic carries an example offending coordinate, the contributing
     PipeNet(s), and a suggested guard.
 
+    Same-thread pipe schedule diagnostics are emitted only when role and DFB
+    domain checks have succeeded, so schedule-cycle reports are not mixed with
+    earlier guard failures.
+
     The pass expects one module per `@ttl.operation`, a well-formed
     `ttl.launch_grid` (i64 array of length 2 with positive entries), static
     `I64Attr` coordinates on `ttl.create_pipe`, and finalized DFB indices

lib/Dialect/TTL/IR/TTLOps.cpp

@@ -176,9 +176,15 @@ mlir::LogicalResult mlir::tt::ttl::CopyOp::verify() {
     if (srcIsPipe && dstIsPipe) {
       return emitOpError() << "cannot copy directly between pipes";
     }
-    if (!srcIsCb && !dstIsCb) {
-      return emitOpError()
-             << "pipe transfers require one operand to be !ttl.cb";
+    if (dstIsPipe) {
+      if (!srcIsCb) {
+        return emitOpError()
+               << "pipe send requires source operand to be !ttl.cb";
+      }
+      return success();
+    }
+    if (!findCBReserveForPipeReceive(getDst())) {
+      return emitOpError() << "pipe receive requires a cb_reserve destination";
     }
     return success();
   }
@@ -247,6 +253,37 @@ mlir::LogicalResult mlir::tt::ttl::CopyOp::verify() {
   return success();
 }
 
+mlir::LogicalResult mlir::tt::ttl::PipeRecvPostOp::verify() {
+  if (!findCBReserveForPipeReceive(getDst())) {
+    return emitOpError() << "requires a cb_reserve destination";
+  }
+
+  auto handleType = mlir::dyn_cast<TransferHandleType>(getXf().getType());
+  if (!handleType) {
+    return emitOpError() << "requires a transfer handle result";
+  }
+  if (handleType.getKind()) {
+    return emitOpError() << "requires an untyped transfer handle result";
+  }
+
+  return success();
+}
+
+mlir::LogicalResult mlir::tt::ttl::PipeRecvWaitOp::verify() {
+  auto handleType = mlir::dyn_cast<TransferHandleType>(getXf().getType());
+  if (!handleType) {
+    return emitOpError() << "requires a transfer handle operand";
+  }
+  if (handleType.getKind()) {
+    return emitOpError() << "requires an untyped transfer handle operand";
+  }
+  if (!findCBReserveForPipeReceive(getDst())) {
+    return emitOpError() << "requires a cb_reserve destination";
+  }
+
+  return success();
+}
+
 mlir::LogicalResult mlir::tt::ttl::WaitOp::verify() {
   if (failed(
           mlir::tt::ttl::verify::isValidWaitOperand(getOperation(), getXf()))) {

lib/Dialect/TTL/IR/TTLOpsVerifyUtils.cpp

@@ -6,93 +6,45 @@
 
 #include "TTLOpsVerifyUtils.h"
 
-#include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Support/LogicalResult.h"
+#include "ttlang/Dialect/TTL/IR/TTLOpsUtils.h"
 #include "llvm/ADT/SmallPtrSet.h"
 
 namespace mlir::tt::ttl::verify {
 namespace {
 
-static bool isDerivedFromCopy(mlir::Value v,
-                              llvm::SmallPtrSetImpl<mlir::Value> &seen) {
-  if (!seen.insert(v).second) {
-    return false;
-  }
-
-  if (v.getDefiningOp<mlir::tt::ttl::CopyOp>() != nullptr) {
-    return true;
-  }
-
-  // tensor.extract: the extracted handle is only valid if the source tensor is
-  // derived from a copy.
-  if (auto extractOp = v.getDefiningOp<mlir::tensor::ExtractOp>()) {
-    return isDerivedFromCopy(extractOp.getTensor(), seen);
-  }
-
-  // tensor.insert: the tensor is derived from a copy if the inserted scalar is
-  // derived from a copy.
-  if (auto insertOp = v.getDefiningOp<mlir::tensor::InsertOp>()) {
-    return isDerivedFromCopy(insertOp.getScalar(), seen);
-  }
-
-  // Loop result -> yielded value.
-  if (auto res = llvm::dyn_cast<mlir::OpResult>(v)) {
-    if (auto loop = llvm::dyn_cast<mlir::LoopLikeOpInterface>(res.getOwner())) {
-      auto yieldedOpt = loop.getYieldedValuesMutable();
-      auto resultsOpt = loop.getLoopResults();
-      if (yieldedOpt && resultsOpt) {
-        auto yielded = *yieldedOpt;
-        auto results = *resultsOpt;
-        for (unsigned idx = 0; idx < results.size(); ++idx) {
-          if (results[idx] != res) {
-            continue;
-          }
-          return isDerivedFromCopy(yielded[idx].get(), seen);
-        }
-      }
-    }
-  }
-
-  // Loop iter_arg -> init.
-  if (auto barg = llvm::dyn_cast<mlir::BlockArgument>(v)) {
-    auto *parent = barg.getOwner()->getParentOp();
-    if (auto loop = llvm::dyn_cast_or_null<mlir::LoopLikeOpInterface>(parent)) {
-      auto iterArgs = loop.getRegionIterArgs();
-      auto inits = loop.getInitsMutable();
-      for (unsigned idx = 0; idx < iterArgs.size(); ++idx) {
-        if (iterArgs[idx] != barg) {
-          continue;
-        }
-        return isDerivedFromCopy(inits[idx].get(), seen);
-      }
-    }
-  }
-
-  return false;
+// Return true when `v` is a transfer handle produced by `ttl.copy`, allowing
+// the handle to flow through tensor containers and loop-carried values.
+// `ttl.pipe_recv_post` is accepted because receive-side `ttl.copy` expands to
+// that internal op before lowering while preserving the original wait contract.
+static bool isDerivedFromCopy(mlir::Value value) {
+  llvm::SmallPtrSet<mlir::Value, 16> seen;
+  return mlir::tt::ttl::traceTransferHandleSource<bool>(
+      value,
+      [](mlir::Value source) {
+        return source.getDefiningOp<mlir::tt::ttl::CopyOp>() != nullptr ||
+               source.getDefiningOp<mlir::tt::ttl::PipeRecvPostOp>() != nullptr;
+      },
+      seen);
 }
 
 } // namespace
 
 mlir::LogicalResult isValidWaitOperand(mlir::Operation *op,
                                        mlir::Value handle) {
-  // Accept any TransferHandleType (typed or untyped).
-  // Typed handles (read/write) get corresponding barriers.
-  // Untyped handles (e.g., pipe receive) are no-ops since data arrives via
-  // multicast from source core.
+  // Accept typed and untyped transfer handles. Untyped handles model async
+  // pipe receive completion and are expanded before lowering.
   if (!mlir::isa<mlir::tt::ttl::TransferHandleType>(handle.getType())) {
     return op->emitOpError()
            << "expects transfer handle (!ttl.transfer_handle), got "
            << handle.getType();
   }
 
-  llvm::SmallPtrSet<mlir::Value, 16> visited;
-  if (isDerivedFromCopy(handle, visited)) {
+  if (isDerivedFromCopy(handle)) {
     return mlir::success();
   }
 
-  return op->emitOpError() << "expects operand to be the result of ttl.copy.";
+  return op->emitOpError() << "expects operand to be derived from ttl.copy.";
 }
 
 } // namespace mlir::tt::ttl::verify

lib/Dialect/TTL/Transforms/CMakeLists.txt

@@ -30,6 +30,7 @@ add_mlir_dialect_library(TTLangTTLTransforms
 
   LINK_LIBS PUBLIC
   MLIRAnalysis
+  MLIRAffineTransforms
   MLIRArithDialect
   MLIREmitCDialect
   MLIRFuncDialect