Fix silent miscompile of thread-varying DMA into shared storage

lib/Target/GPU/cute_codegen.cpp

@@ -2754,7 +2754,13 @@ bool CuteCodeGen::Visit(AST::NamedVariableDecl& n) {
         auto reuse = n.GetNote("reuse");
         auto offset = n.GetNote("offset");
         ds << d_indent << bts << "* " << sym << " = (" << bts << "*)"
-           << "(" << reuse << " + " << offset << ");\n";
+           << "(" << reuse << " + " << offset;
+        // Thread-sliced shared buffer: each thread stages its own slice of
+        // the spm, offset by vtid * (per-thread slice bytes).
+        if (n.HasNote("tsliced"))
+          ds << " + __choreo_vtid_x * (" << UnScopedExpr(n.GetNote("tsliced"))
+             << ")";
+        ds << ");\n";
       } else {
         // the buffer is not reused
         assert(!n.HasNote("offset"));
@@ -2874,7 +2880,10 @@ bool CuteCodeGen::Visit(AST::NamedVariableDecl& n) {
       if (sto != Storage::SHARED && sto != Storage::LOCAL)
         choreo_unreachable(
             "error: unexpected storage type in spm initialization.");
-      if (sto == Storage::SHARED) {
+      // A thread-sliced shared buffer is initialized per-thread (each thread
+      // zeroes its own slice); block-common buffers are initialized once.
+      bool tsliced = n.HasNote("tsliced");
+      if (sto == Storage::SHARED && !tsliced) {
         ds << d_indent << LevelPred() << " {\n";
         IncrDeviceIndent();
       }
@@ -2886,8 +2895,10 @@ bool CuteCodeGen::Visit(AST::NamedVariableDecl& n) {
         ds << sym << "[i] = " << ExprSTR(n.init_value) << ";\n";
       }
       if (sto == Storage::SHARED) {
-        DecrDeviceIndent();
-        ds << d_indent << "} // single instance\n";
+        if (!tsliced) {
+          DecrDeviceIndent();
+          ds << d_indent << "} // single instance\n";
+        }
         ds << d_indent << EmitSync(Storage::SHARED) << ";\n";
       }
     }
@@ -4220,7 +4231,12 @@ bool CuteCodeGen::Visit(AST::DMA& n) {
       bool to_or_from_shared = dma_plan->direction == DMADirection::G2S ||
                                dma_plan->direction == DMADirection::S2G ||
                                dma_plan->direction == DMADirection::S2S;
-      if (to_or_from_shared) {
+      if (to_or_from_shared && dma_plan->thread_sliced) {
+        // Thread-sliced shared buffer: every thread copies its own slice;
+        // a BLOCK_SINGLE guard would move only thread 0's slice.
+        ds << d_indent << "choreo::naive_copy(" << src << ", " << dst << ");\n";
+        ds << d_indent << "__syncthreads();\n";
+      } else if (to_or_from_shared) {
         ds << d_indent << "if (__CHOREO_BLOCK_SINGLE__) {\n";
         ds << d_indent << "  choreo::naive_copy(" << src << ", " << dst
            << ");\n";
@@ -10002,17 +10018,26 @@ void CuteCodeGen::EmitHostRuntimeCheck() {
     for (size_t i = 1; i < entries.size(); ++i) {
       auto& entry0 = entries[i - 1];
       auto& entry1 = entries[i];
+      std::string msg = "The shapes of the " + Ordinal(entry0.para_ordinal) +
+                        " parameter (dim: " + std::to_string(entry0.dim) +
+                        ") and the " + Ordinal(entry1.para_ordinal) +
+                        " parameter (dim: " + std::to_string(entry1.dim) +
+                        ") are inconsistent.";
       hs << h_indent << "choreo::runtime_check(" << entry0.elem_name
          << " == " << entry1.elem_name;
-      hs << ", \"The shapes of the " << Ordinal(entry0.para_ordinal)
-         << " parameter (dim: " << entry0.dim << ") and the "
-         << Ordinal(entry1.para_ordinal) << " parameter (dim: " << entry1.dim
-         << ") are inconsistent.\");\n";
+      hs << ", \"" << msg << "\");\n";
       ++stats.total;
       ++stats.shape_compat_total;
       ++stats.runtime_total;
       ++stats.shape_compat_runtime;
       ++stats.runtime_entry;
+      // Record in the safety ledger: these obligations are generated and
+      // materialized directly at codegen; parameter dims are host-visible
+      // scalars, hence scalar-symbolic dependence. Stats aggregation already
+      // happened (assert-site pass), so this does not double-count.
+      FCtx(fname).GetAssessor().RecordExternalObligation(
+          msg, location(), AssessOutcome::RUNTIME,
+          UsageType::ShapeCompatibility, AssessDependence::SCALAR_SYMBOLIC);
     }
   }
 

lib/Target/GPU/dma_plan.cpp

@@ -1,5 +1,6 @@
 #include "dma_plan.hpp"
 #include "ast.hpp"
+#include "thread_sliced_shared.hpp"
 #include "types.hpp"
 
 using namespace Choreo;
@@ -234,6 +235,20 @@ void DMAPlan::ResolveDMADecision(const AST::DMA& n,
   dec.rank = static_cast<int>(from_ca->GetBlockShape().Rank());
   dec.elem_type = from_sty->ElementType();
 
+  // A shared buffer staged per-thread (thread-sliced, see
+  // thread_sliced_shared.hpp) must be copied per-thread: cooperative tiled
+  // lowering would fill only one slice, and a __CHOREO_BLOCK_SINGLE__ guard
+  // would load only thread 0's slice.  Mark the decision so codegen emits an
+  // unguarded per-thread copy; the tiled-copy ladder below is skipped.
+  auto touches_tsliced = [this](const std::string& sym) {
+    return ThreadSlicedShared::Lookup(InScopeName(sym)) != nullptr;
+  };
+  if ((dec.direction == DMADirection::G2S && touches_tsliced(to_sym)) ||
+      (dec.direction == DMADirection::S2G && touches_tsliced(from_sym)) ||
+      (dec.direction == DMADirection::S2S &&
+       (touches_tsliced(from_sym) || touches_tsliced(to_sym))))
+    dec.thread_sliced = true;
+
   if (n.IsTMA()) {
     dec.strategy = DMAStrategy::TMA;
     dec.atom = CUDA_COPY_ATOM::TMA_ATOM;
@@ -327,6 +342,11 @@ void DMAPlan::ResolveDMADecision(const AST::DMA& n,
     debug_naive_fallback("chunk indexing depends on inner-than-block PV");
     return;
   }
+  if (dec.thread_sliced) {
+    debug_naive_fallback(
+        "thread-sliced shared buffer requires per-thread copy");
+    return;
+  }
 
   if (n.operation != ".copy" && n.operation != ".transp" &&
       n.operation != ".pad") {

lib/Target/GPU/dma_plan.hpp

@@ -183,6 +183,10 @@ struct DMALoweringDecision {
   bool is_zfill = false; // dma.copy.*.zfill  -> cp.async zero-fill semantic
   bool has_pred = false; // predicated tiled copy needed (tail handling)
   bool use_tma = false;  // TMA bulk-copy path
+  // The DMA moves per-thread slices of a thread-sliced shared buffer; it
+  // must be lowered to a per-thread naive copy (no __CHOREO_BLOCK_SINGLE__
+  // guard, no cooperative tiled copy).
+  bool thread_sliced = false;
 
   // -- swizzle ---------------------------------------------------------------
   SwizMode swizzle_mode = SwizMode::NONE;

lib/assert_site.cpp

@@ -493,10 +493,26 @@ void AssertSite::EstimateAssertions() {
       case UsageType::HardwareConstraint: ++stats.hw_constraint_total; break;
       }
       switch (ae.outcome) {
-      case AssessOutcome::STATIC_TRUE: ++stats.static_true; break;
+      case AssessOutcome::STATIC_TRUE:
+        ++stats.static_true;
+        switch (ae.dependence) {
+        case AssessDependence::STRUCTURAL: ++stats.st_structural; break;
+        case AssessDependence::SCALAR_SYMBOLIC: ++stats.st_scalar; break;
+        case AssessDependence::CONSTANT: ++stats.st_const; break;
+        }
+        switch (ae.mechanism) {
+        case AssessMechanism::CANONICAL: ++stats.st_canonical; break;
+        case AssessMechanism::INTERVAL: ++stats.st_interval; break;
+        }
+        break;
       case AssessOutcome::STATIC_FALSE: ++stats.static_false; break;
       case AssessOutcome::RUNTIME: {
         ++stats.runtime_total;
+        switch (ae.dependence) {
+        case AssessDependence::STRUCTURAL: ++stats.rt_structural; break;
+        case AssessDependence::SCALAR_SYMBOLIC: ++stats.rt_scalar; break;
+        case AssessDependence::CONSTANT: ++stats.rt_const; break;
+        }
         // Per-usage-type runtime
         switch (ae.usage_type) {
         case UsageType::UnClassified: ++stats.unclassified_runtime; break;

lib/assess.cpp

@@ -47,32 +47,13 @@ inline const std::string STR(const AssessRelation& ar) {
   return "";
 }
 
-inline const std::string STR(const AssessOutcome& o) {
-  switch (o) {
-  case AssessOutcome::STATIC_TRUE: return "static-true";
-  case AssessOutcome::STATIC_FALSE: return "static-false";
-  case AssessOutcome::RUNTIME: return "runtime";
-  }
-  return "?";
-}
-
-inline const std::string STR(const UsageType& ut) {
-  switch (ut) {
-  case UsageType::UnClassified: return "unclassified";
-  case UsageType::ShapeCompatibility: return "shape-compat";
-  case UsageType::ElementAccess: return "elem-access";
-  case UsageType::LoopBound: return "loop-bound";
-  case UsageType::HardwareConstraint: return "hw-constraint";
-  }
-  choreo_unreachable("unsupported usage type.");
-  return "";
-}
 } // namespace Choreo
 
 void Assessor::LogAssessment(const std::string& msg, const location& l,
                              AssessOutcome outcome, UsageType uty,
-                             size_t assertion_idx) {
-  assessment_log.push_back({msg, l, outcome, uty, assertion_idx});
+                             AssessDependence dep, size_t assertion_idx,
+                             AssessMechanism mech) {
+  assessment_log.push_back({msg, l, outcome, uty, dep, mech, assertion_idx});
 }
 
 void Assessor::AddAssertion(const ptr<sbe::SymbolicExpression>& ar,
@@ -115,22 +96,23 @@ AssessResult Assessor::Assess(AssessPolicy ap, AssessRelation rel,
   auto pred =
       (rel == AssessRelation::EQ) ? sbe::oc_eq(lhs, rhs) : sbe::oc_ne(lhs, rhs);
   const auto warning_msg = warn_message.empty() ? error_message : warn_message;
+  const auto dep = ClassifyDependence({&lhs, &rhs});
 
   if (auto b = VIBool(pred)) {
     if (b.value() == false) {
       switch (ap) {
       case AssessPolicy::Error:
       case AssessPolicy::ErrWarn:
         visitor->Error1(l, error_message);
-        LogAssessment(error_message, l, AssessOutcome::STATIC_FALSE, uty);
+        LogAssessment(error_message, l, AssessOutcome::STATIC_FALSE, uty, dep);
         return {false, false, false};
       case AssessPolicy::Warn:
         visitor->Warning(l, warning_msg);
-        LogAssessment(warning_msg, l, AssessOutcome::STATIC_FALSE, uty);
+        LogAssessment(warning_msg, l, AssessOutcome::STATIC_FALSE, uty, dep);
         return {true, true, false};
       }
     }
-    LogAssessment(error_message, l, AssessOutcome::STATIC_TRUE, uty);
+    LogAssessment(error_message, l, AssessOutcome::STATIC_TRUE, uty, dep);
     return {true, false, false};
   }
 
@@ -148,7 +130,7 @@ AssessResult Assessor::Assess(AssessPolicy ap, AssessRelation rel,
   case AssessPolicy::Error:
     if (strict_fail) {
       visitor->Error1(l, error_message);
-      LogAssessment(error_message, l, AssessOutcome::STATIC_FALSE, uty);
+      LogAssessment(error_message, l, AssessOutcome::STATIC_FALSE, uty, dep);
       return {false, false, false};
     }
     break;
@@ -160,19 +142,19 @@ AssessResult Assessor::Assess(AssessPolicy ap, AssessRelation rel,
     LogAssessment(error_message, l,
                   strict_fail ? AssessOutcome::STATIC_FALSE
                               : AssessOutcome::STATIC_TRUE,
-                  uty);
+                  uty, dep);
     return {true, strict_fail || may_fail, false};
   case AssessPolicy::ErrWarn:
     if (strict_fail) {
       visitor->Error1(l, error_message);
-      LogAssessment(error_message, l, AssessOutcome::STATIC_FALSE, uty);
+      LogAssessment(error_message, l, AssessOutcome::STATIC_FALSE, uty, dep);
       return {false, false, false};
     }
     if (may_fail) visitor->Warning(l, warning_msg);
     break;
   }
 
-  LogAssessment(error_message, l, AssessOutcome::RUNTIME, uty,
+  LogAssessment(error_message, l, AssessOutcome::RUNTIME, uty, dep,
                 assertions.size());
   AddAssertion(pred, l, error_message, aty, uty, node);
   return {true, may_fail, true};
@@ -190,7 +172,9 @@ AssessResult Assessor::Assess(AssessPolicy ap, const ValueItem& bo,
                               const std::string& message, UsageType uty,
                               AssessType aty, const location& l,
                               AST::Node* node, AST::Node* emit_node,
-                              const ValueItem& guard) {
+                              const ValueItem& guard,
+                              std::optional<AssessDependence> dep_override,
+                              AssessMechanism mech) {
   if (DebugOn())
     dbgs() << "[Assess] " << STR(bo) << ", type: " << STR(aty)
            << ", usage: " << STR(uty) << ", policy: " << STR(ap)
@@ -205,6 +189,7 @@ AssessResult Assessor::Assess(AssessPolicy ap, const ValueItem& bo,
 
   auto pred = bo;
   if (pred) pred = pred->Normalize();
+  const auto dep = dep_override.value_or(ClassifyDependence({&bo}));
 
   auto norm_guard = guard;
   if (norm_guard) norm_guard = norm_guard->Normalize();
@@ -219,25 +204,28 @@ AssessResult Assessor::Assess(AssessPolicy ap, const ValueItem& bo,
       if (IsValidValueItem(norm_guard)) {
         // Statically false but only reachable under a guard; keep as runtime.
         if (ap == AssessPolicy::Warn) return {true, false, false};
-        LogAssessment(message, l, AssessOutcome::RUNTIME, uty,
-                      assertions.size());
+        LogAssessment(message, l, AssessOutcome::RUNTIME, uty, dep,
+                      assertions.size(), mech);
         AddAssertion(pred, l, message, aty, uty, node, emit_node);
         return {true, false, true};
       }
       if (ap == AssessPolicy::Error)
         visitor->Error1(l, message);
       else
         visitor->Warning(l, message);
-      LogAssessment(message, l, AssessOutcome::STATIC_FALSE, uty);
+      LogAssessment(message, l, AssessOutcome::STATIC_FALSE, uty, dep,
+                    static_cast<size_t>(-1), mech);
       return {ap == AssessPolicy::Warn, ap == AssessPolicy::Warn, false};
     }
-    LogAssessment(message, l, AssessOutcome::STATIC_TRUE, uty);
+    LogAssessment(message, l, AssessOutcome::STATIC_TRUE, uty, dep,
+                  static_cast<size_t>(-1), mech);
     return {true, false, false};
   }
 
   if (ap == AssessPolicy::Warn) return {true, false, false};
 
-  LogAssessment(message, l, AssessOutcome::RUNTIME, uty, assertions.size());
+  LogAssessment(message, l, AssessOutcome::RUNTIME, uty, dep,
+                assertions.size(), mech);
   AddAssertion(pred, l, message, aty, uty, node, emit_node);
   return {true, false, true};
 }