diff --git a/lib/Target/GPU/cute_codegen.cpp b/lib/Target/GPU/cute_codegen.cpp index 04fa7bca..c755249f 100644 --- a/lib/Target/GPU/cute_codegen.cpp +++ b/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); } } diff --git a/lib/Target/GPU/dma_plan.cpp b/lib/Target/GPU/dma_plan.cpp index 200c15dd..8f9b5e44 100644 --- a/lib/Target/GPU/dma_plan.cpp +++ b/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(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") { diff --git a/lib/Target/GPU/dma_plan.hpp b/lib/Target/GPU/dma_plan.hpp index 1d8ab328..34fdbaf4 100644 --- a/lib/Target/GPU/dma_plan.hpp +++ b/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; diff --git a/lib/assert_site.cpp b/lib/assert_site.cpp index e26d76a1..a3e94dce 100644 --- a/lib/assert_site.cpp +++ b/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; diff --git a/lib/assess.cpp b/lib/assess.cpp index d0f06f61..53f076cb 100644 --- a/lib/assess.cpp +++ b/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& ar, @@ -115,6 +96,7 @@ 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) { @@ -122,15 +104,15 @@ AssessResult Assessor::Assess(AssessPolicy ap, AssessRelation rel, 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 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,8 +204,8 @@ 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}; } @@ -228,16 +213,19 @@ AssessResult Assessor::Assess(AssessPolicy ap, const ValueItem& bo, 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(-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(-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}; } diff --git a/lib/assess.hpp b/lib/assess.hpp index 2184387f..d09ec1ea 100644 --- a/lib/assess.hpp +++ b/lib/assess.hpp @@ -4,6 +4,8 @@ #include "loc.hpp" #include "symvals.hpp" #include +#include +#include #include #include @@ -39,6 +41,18 @@ enum class AssertionEmitPosition { IN_BLOCK, }; +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 ""; +} + enum class AssertionCost { NONE, ENTRY, @@ -65,12 +79,43 @@ enum class AssessOutcome { RUNTIME, ///< Cannot evaluate -- runtime assertion emitted. }; +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 "?"; +} + +/// Information dependence of an assessment's predicate (RQ4 capability-gap +/// analysis): what class of information the discharge relies on. +enum class AssessDependence { + CONSTANT, ///< No symbolic content (constants only). + SCALAR_SYMBOLIC, ///< Plain scalar symbols (params/dims) that survive + ///< lowering as runtime values. + STRUCTURAL, ///< References block/thread-parallel structure (symbols + ///< scoped under paraby_/inthreads_) that lowering + ///< dissolves into flat index arithmetic. +}; + +/// Resolution mechanism that decided an assessment (RQ6 mechanism +/// breakdown). CANONICAL: structural normalization/constant folding of the +/// predicate itself. INTERVAL: interval analysis over bounded-type ranges +/// (TryProveWithIntervals in the semantic checker). +enum class AssessMechanism { + CANONICAL, + INTERVAL, +}; + /// Record of every assessment evaluation, regardless of outcome. struct AssessmentEntry { std::string message; location loc; AssessOutcome outcome; UsageType usage_type = UsageType::UnClassified; + AssessDependence dependence = AssessDependence::CONSTANT; + AssessMechanism mechanism = AssessMechanism::CANONICAL; /// Index into Assessor::assertions (RUNTIME only); SIZE_MAX otherwise. size_t assertion_idx = static_cast(-1); }; @@ -81,6 +126,28 @@ struct AssessResult { bool inserted = false; }; +// Classify the information dependence of an assessment predicate from the +// symbols referenced by its (pre-decision) operands: structural symbols are +// those scoped under block/thread-parallel boundaries (paraby_/inthreads_), +// which exist only at the semantic stage. +inline AssessDependence ClassifyDependence( + std::initializer_list operands) { + bool has_symbol = false; + for (auto* op : operands) { + if (!op || !IsValidValueItem(*op)) continue; + for (const auto& sym : GetSymbols(*op)) { + has_symbol = true; + if (auto name = VISym(sym)) { + if (name->find("paraby_") != std::string::npos || + name->find("inthreads_") != std::string::npos) + return AssessDependence::STRUCTURAL; + } + } + } + return has_symbol ? AssessDependence::SCALAR_SYMBOLIC + : AssessDependence::CONSTANT; +} + struct Assertion { ptr expr; @@ -118,7 +185,9 @@ class Assessor { /// Record a single assessment evaluation to the ordered log. void LogAssessment(const std::string& msg, const location& l, AssessOutcome outcome, UsageType uty, - size_t assertion_idx = static_cast(-1)); + AssessDependence dep = AssessDependence::CONSTANT, + size_t assertion_idx = static_cast(-1), + AssessMechanism mech = AssessMechanism::CANONICAL); bool DebugOn() const; @@ -136,6 +205,16 @@ class Assessor { return assessment_log; } + /// Record an obligation that was assessed/materialized outside the Assess() + /// path (e.g. parameter shape-consistency checks emitted at codegen) so the + /// safety ledger remains complete. Does not affect stats aggregation, which + /// has already replayed the log by codegen time. + void RecordExternalObligation(const std::string& msg, const location& l, + AssessOutcome outcome, UsageType uty, + AssessDependence dep) { + LogAssessment(msg, l, outcome, uty, dep); + } + std::vector GetAssertions(AssessType aty) const { std::vector output; output.reserve(assertions.size()); @@ -162,7 +241,10 @@ class Assessor { const std::string& message, UsageType uty, AssessType aty, const location& l, AST::Node* node, AST::Node* emit_node = nullptr, - const ValueItem& guard = GetInvalidValueItem()); + const ValueItem& guard = GetInvalidValueItem(), + std::optional dep_override = + std::nullopt, + AssessMechanism mech = AssessMechanism::CANONICAL); }; } // end namespace Choreo diff --git a/lib/command_line.cpp b/lib/command_line.cpp index 1f470562..8f4ad044 100644 --- a/lib/command_line.cpp +++ b/lib/command_line.cpp @@ -211,8 +211,20 @@ Option dump_ast(OptionKind::User, "--dump-ast", "-e", false, "Dump the Abstract Syntax Tree (AST) after parsing."); Option dump_hb(OptionKind::User, "--dump-hb", "", false, "Dump Happens-Before graph in DOT (Graphviz) format."); +Option dump_ledger( + OptionKind::Hidden, "--dump-ledger", "", "" /*default empty*/, + "Dump the safety ledger (all assessed obligations with outcomes) as JSON."); Option print_vn(OptionKind::Hidden, "--print-valno", "-vn", false, "Trace the value numbering process."); +Option disable_vn_share( + OptionKind::Hidden, "--disable-vn-share", "", false, + "(Experimental) Disable value-number sharing while keeping " + "canonicalization, for ablation studies."); +Option disable_vn_simplify( + OptionKind::Hidden, "--disable-vn-simplify", "", false, + "(Experimental) Disable value-number-based simplification rules " + "(keeping constant folding and structural normalization), for " + "ablation studies."); Option dump_sym(OptionKind::Hidden, "--dump-symbol", "-l", false, "Dump the symbol table after LATENORM."); Option visualiz(OptionKind::Hidden, "--visualize", "-u", false, @@ -467,6 +479,7 @@ bool CommandLine::Parse(int argc, char** argv) { target_generate_debug_info.GetValue()); CCtx().SetDumpAst(dump_ast.GetValue()); CCtx().SetDumpHB(dump_hb.GetValue()); + CCtx().SetDumpLedgerPath(dump_ledger.GetValue()); CCtx().SetNoCodegen(ncodegen.GetValue()); CCtx().SetPrintPassNames(prt_pass.GetValue()); CCtx().SetTimePasses(time_passes.GetValue()); @@ -479,6 +492,8 @@ bool CommandLine::Parse(int argc, char** argv) { CCtx().SetVisualize(visualiz.GetValue()); CCtx().SetCrossCompile(cross_compile.GetValue()); CCtx().SetTraceValueNumbers(print_vn.GetValue()); + CCtx().SetDisableVNShare(disable_vn_share.GetValue()); + CCtx().SetDisableVNSimplify(disable_vn_simplify.GetValue()); CCtx().SetTraceVectorize(debug_vectorize.GetValue()); CCtx().SetNoVectorize(no_vectorize.GetValue()); CCtx().SetNoMapHoist(no_map_hoist.GetValue()); diff --git a/lib/context.hpp b/lib/context.hpp index 28d91ee6..90a7716f 100644 --- a/lib/context.hpp +++ b/lib/context.hpp @@ -469,6 +469,17 @@ struct AssessmentStats { size_t static_true = 0; // resolved at compile time (always passes) size_t static_false = 0; // proven false at compile time (error/warning) size_t runtime_total = 0; // runtime assertions generated + // Discharge information dependence (Assessor-logged path only; RQ4) + size_t st_structural = 0; // static_true, needs block/thread structure + size_t st_scalar = 0; // static_true, scalar symbols only + size_t st_const = 0; // static_true, constants only + size_t rt_structural = 0; // runtime, references block/thread structure + size_t rt_scalar = 0; // runtime, scalar symbols only + size_t rt_const = 0; // runtime, constants only + size_t direct_checks = 0; // assessed via direct static checks (no log) + // Discharge resolution mechanism (Assessor-logged path; RQ6) + size_t st_canonical = 0; // static_true via canonical normalization + size_t st_interval = 0; // static_true via interval/bounded-type proof size_t runtime_entry = 0; // runtime assertions with entry estimated cost size_t runtime_low = 0; // runtime assertions with low estimated cost size_t runtime_medium = 0; // runtime assertions with medium estimated cost @@ -487,7 +498,6 @@ struct AssessmentStats { size_t elem_access_runtime = 0; size_t loop_bound_runtime = 0; size_t hw_constraint_runtime = 0; - // Automatic DMA fence-insertion statistics (gated by --collect-stats). struct FenceStats { size_t inserted = 0; // auto fences emitted (producer + consumer) @@ -536,6 +546,7 @@ class CompilationContext { bool debug_symtab = false; bool dump_ast = false; // dump the AST after parsing bool dump_hb = false; // dump HB graph in DOT format + std::string dump_ledger_path; // dump the safety ledger as JSON bool no_codegen = false; // stop before code generation bool print_pass_names = false; // print pass name before pass run bool time_passes = false; // measure time per compiler pass @@ -548,6 +559,8 @@ class CompilationContext { bool visualize = false; // visualize the DMAs bool cross_compile = false; // TODO: figure out bool trace_vn = false; // trace the value numbering + bool disable_vn_share = false; // disable value-number sharing (ablation) + bool disable_vn_simplify = false; // disable VN-based simplification (ablation) bool trace_vectorize = false; // trace the masking bool show_source_loc = true; // show source code location when error, etc. bool mem_reuse = false; // reuse the memory of the program @@ -810,6 +823,7 @@ class CompilationContext { // Getters of compiler configurations bool DumpAst() const { return dump_ast; } bool DumpHB() const { return dump_hb; } + const std::string& DumpLedgerPath() const { return dump_ledger_path; } bool NoCodegen() const { return no_codegen; } bool PrintPassNames() const { return print_pass_names; } bool TimePasses() const { return time_passes; } @@ -822,6 +836,8 @@ class CompilationContext { bool Visualize() const { return visualize; } bool CrossCompile() const { return cross_compile; } bool TraceValueNumbers() const { return trace_vn; } + bool DisableVNShare() const { return disable_vn_share; } + bool DisableVNSimplify() const { return disable_vn_simplify; } bool TraceVectorize() const { return trace_vectorize; } bool MemReuse() const { return mem_reuse; } bool SALA() const { return sala; } @@ -877,6 +893,7 @@ class CompilationContext { // Setters of compiler configurations void SetDumpAst(bool value) { dump_ast = value; } void SetDumpHB(bool value) { dump_hb = value; } + void SetDumpLedgerPath(const std::string& path) { dump_ledger_path = path; } void SetNoCodegen(bool value) { no_codegen = value; } void SetPrintPassNames(bool value) { print_pass_names = value; } void SetTimePasses(bool value) { time_passes = value; } @@ -889,6 +906,8 @@ class CompilationContext { void SetVisualize(bool value) { visualize = value; } void SetCrossCompile(bool value) { cross_compile = value; } void SetTraceValueNumbers(bool value) { trace_vn = value; } + void SetDisableVNShare(bool value) { disable_vn_share = value; } + void SetDisableVNSimplify(bool value) { disable_vn_simplify = value; } void SetTraceVectorize(bool value) { trace_vectorize = value; } void SetMemReuse(bool value) { mem_reuse = value; } void SetSALA(bool value) { sala = value; } diff --git a/lib/mem_reuse.cpp b/lib/mem_reuse.cpp index 629975c1..a071895f 100644 --- a/lib/mem_reuse.cpp +++ b/lib/mem_reuse.cpp @@ -198,6 +198,17 @@ bool MemAnalyzer::Visit(AST::NamedVariableDecl& n) { buf_size.emplace(sname, size_expr); VST_DEBUG(dbgs() << "\tdynamic size: " << size_expr << "\n";); } + // A thread-sliced shared buffer holds one slice per thread: scale its + // footprint by the slice count and remember the per-thread slice bytes + // so codegen can offset each thread into its own slice. + if (sto == Storage::SHARED) + if (auto count = ThreadSlicedShared::Lookup(sname)) { + buf_tslice_bytes.emplace(sname, buf_size.at(sname)); + buf_tslice_count.emplace(sname, *count); + buf_size[sname] = buf_size.at(sname) * (*count); + VST_DEBUG(dbgs() << "\tthread-sliced: x" << *count << " -> " + << buf_size.at(sname) << "\n";); + } buf_dev_func_name.emplace(sname, cur_dev_fname); VST_DEBUG(dbgs() << "\tdecl in dev func: " << cur_dev_fname << "\n";); return true; @@ -834,6 +845,10 @@ void MemReuse::ApplyMemOffset(AST::NamedVariableDecl& n, Storage sto) { n.AddNote("offset", offset); size_t alignment = AlignmentForDevFunc(sto, cur_dev_fname); n.AddNote("alignment", std::to_string(alignment)); + // Thread-sliced shared buffer: record the per-thread slice bytes so the + // codegen can offset each thread into its own slice of the spm. + if (ma.buf_tslice_bytes.count(sname)) + n.AddNote("tsliced", STR(ma.buf_tslice_bytes.at(sname))); } bool MemReuse::RunOnProgramImpl(AST::Node& root) { diff --git a/lib/mem_reuse.hpp b/lib/mem_reuse.hpp index b56f8a57..b2dc3b25 100644 --- a/lib/mem_reuse.hpp +++ b/lib/mem_reuse.hpp @@ -7,6 +7,7 @@ #include "heap_simulator.hpp" #include "liveness_analysis.hpp" #include "symvals.hpp" +#include "thread_sliced_shared.hpp" #include "typeresolve.hpp" #include "types.hpp" #include "visitor.hpp" @@ -33,6 +34,10 @@ struct MemAnalyzer : public VisitorWithSymTab { std::unordered_map buf_dev_func_name; std::unordered_map buf_alignment; std::set event_vars; + // Thread-sliced shared buffers: scoped name -> per-thread slice bytes + // (footprint before scaling by the thread count) and the slice count. + std::unordered_map buf_tslice_bytes; + std::unordered_map buf_tslice_count; MemAnalyzer() : VisitorWithSymTab("memanlz") {} ~MemAnalyzer() {} diff --git a/lib/pipeline.cpp b/lib/pipeline.cpp index fa20d983..3228a4b1 100644 --- a/lib/pipeline.cpp +++ b/lib/pipeline.cpp @@ -18,9 +18,11 @@ #include "shapeinfer.hpp" #include "symbexpr.hpp" #include "target_utils.hpp" +#include "thread_sliced_shared.hpp" #include "typeinfer.hpp" #include "visualize.hpp" #include +#include #include extern Choreo::AST::Program root; @@ -118,6 +120,94 @@ void ASTPipeline::PrintPassTimings(const std::vector& timings, void ASTPipeline::ValidatePassNames() const { Visitor::ValidatePassEnvVars(); } +namespace { +// --- Safety ledger dump (--dump-ledger) --- + +const char* LedgerDepName(Choreo::AssessDependence d) { + using Choreo::AssessDependence; + switch (d) { + case AssessDependence::CONSTANT: return "constant"; + case AssessDependence::SCALAR_SYMBOLIC: return "scalar-symbolic"; + case AssessDependence::STRUCTURAL: return "structural"; + } + return "?"; +} + +const char* LedgerMechName(Choreo::AssessMechanism m) { + using Choreo::AssessMechanism; + switch (m) { + case AssessMechanism::CANONICAL: return "canonical"; + case AssessMechanism::INTERVAL: return "interval"; + } + return "?"; +} + +const char* LedgerCostName(Choreo::AssertionCost c) { + using Choreo::AssertionCost; + switch (c) { + case AssertionCost::ENTRY: return "entry"; + case AssertionCost::LOW: return "low"; + case AssertionCost::MEDIUM: return "medium"; + case AssertionCost::HIGH: return "high"; + case AssertionCost::NONE: return "none"; + } + return "?"; +} + +std::string LedgerEscape(const std::string& s) { + std::string out; + out.reserve(s.size()); + for (char c : s) { + switch (c) { + case '"': out += "\\\""; break; + case '\\': out += "\\\\"; break; + case '\n': out += "\\n"; break; + case '\t': out += "\\t"; break; + default: out += c; + } + } + return out; +} + +// Dump every assessed obligation (the safety ledger) as JSON. +void DumpAssessmentLedger(const std::string& path) { + using namespace Choreo; + std::ofstream out(path); + if (!out) { + errs() << "error: cannot open ledger output '" << path << "'\n"; + return; + } + out << "{\n \"obligations\": [\n"; + bool first = true; + for (const auto& [fname, fctx] : CCtx().GetAllFunctionContexts()) { + const auto& log = fctx.GetAssessor().GetAssessmentLog(); + const auto& assertions = fctx.GetAssessor().GetAssertions(); + for (const auto& e : log) { + if (!first) out << ",\n"; + first = false; + const auto& pos = e.loc.begin; + out << " {\"function\": \"" << LedgerEscape(fname) << "\"" + << ", \"loc\": \"" << LedgerEscape(pos.filename) << ":" << pos.line + << "." << pos.column << "\"" + << ", \"message\": \"" << LedgerEscape(e.message) << "\"" + << ", \"outcome\": \"" << STR(e.outcome) << "\"" + << ", \"usage\": \"" << STR(e.usage_type) << "\"" + << ", \"dependence\": \"" << LedgerDepName(e.dependence) << "\"" + << ", \"mechanism\": \"" << LedgerMechName(e.mechanism) << "\""; + if (e.outcome == AssessOutcome::RUNTIME && + e.assertion_idx < assertions.size()) { + const auto& a = assertions[e.assertion_idx]; + out << ", \"cost\": \"" << LedgerCostName(a.cost) << "\"" + << ", \"enabled\": " << (a.enabled ? "true" : "false"); + } + out << "}"; + } + } + out << "\n ]\n}\n"; +} + +} // end anonymous namespace + bool ASTPipeline::RunOnProgram(AST::Node& root) { if (debug) Dump(); // verify the input @@ -259,6 +349,9 @@ bool ASTPipeline::RunOnProgram(AST::Node& root) { << "\n"; } + if (!CCtx().DumpLedgerPath().empty()) + DumpAssessmentLedger(CCtx().DumpLedgerPath()); + return !abend; } @@ -310,8 +403,10 @@ ASTPipeline& ASTPipeline::PlanSemanticRoutine() { [](ASTPipeline& p) { CCtx().SetGlobalSymbolTable(p.LastSymTab()); }); } - if (CCtx().TargetSupportMemAlloc() && CCtx().MemReuse()) + if (CCtx().TargetSupportMemAlloc() && CCtx().MemReuse()) { + AddStage(); AddStage(); + } // compute active thread counts for inthreads scopes (before semacheck // so the checker can validate events with thread count info) @@ -373,6 +468,17 @@ void Choreo::PrintAssessmentStats(const AssessmentStats& s) { row(s.runtime_high, "Runtime assertions (high cost)"); row(s.runtime_enabled, "Runtime assertions enabled"); row(s.runtime_disabled, "Runtime assertions disabled (cost or duplicate)"); + errs() << color::err(color::kDim) << " ---" << color::err(color::kReset) + << "\n"; + row(s.st_structural, "Static-true needing block/thread structure"); + row(s.st_scalar, "Static-true over scalar symbols only"); + row(s.st_const, "Static-true constant-only"); + row(s.rt_structural, "Runtime assertions referencing block/thread structure"); + row(s.rt_scalar, "Runtime assertions over scalar symbols only"); + row(s.rt_const, "Runtime assertions constant-only"); + row(s.direct_checks, "Direct static checks (bypassing assessor)"); + row(s.st_canonical, "Static-true via canonical normalization"); + row(s.st_interval, "Static-true via interval/bounded-type proof"); errs() << color::err(color::kDim) << " ---" << color::err(color::kReset) << "\n"; row(s.unclassified_total, "Assessments (unclassified)"); diff --git a/lib/semacheck.cpp b/lib/semacheck.cpp index f73bf057..2d5b3abe 100644 --- a/lib/semacheck.cpp +++ b/lib/semacheck.cpp @@ -2111,12 +2111,17 @@ void SemaChecker::CreateAssessment(const ValueItem& pred, // Try to prove the predicate using interval analysis on variable ranges // narrowed by both BoundedType declarations and active scope predicates. auto effective_pred = pred; + auto mech = AssessMechanism::CANONICAL; if (!VIBool(pred) && IsValidValueItem(active_guard)) { auto proven = TryProveWithIntervals(this, pred, active_guard); - if (proven.has_value()) effective_pred = sbe::bl(*proven); + if (proven.has_value()) { + effective_pred = sbe::bl(*proven); + mech = AssessMechanism::INTERVAL; + } } FCtx(fname).GetAssessor(*this).Assess(AssessPolicy::Error, effective_pred, message, uty, aty, l, n.get(), - emit_node, active_guard); + emit_node, active_guard, + ClassifyDependence({&pred}), mech); } diff --git a/lib/shapeinfer.cpp b/lib/shapeinfer.cpp index d9929df1..f8b6f879 100644 --- a/lib/shapeinfer.cpp +++ b/lib/shapeinfer.cpp @@ -37,6 +37,7 @@ void ShapeInference::InvalidateVisitorValNOs() { bool ShapeInference::StaticFail(bool pred_fail, UsageType ut) { auto& stats = CCtx().GetAssessmentStats(); ++stats.total; + ++stats.direct_checks; if (pred_fail) ++stats.static_false; else diff --git a/lib/thread_sliced_shared.hpp b/lib/thread_sliced_shared.hpp new file mode 100644 index 00000000..e497e6c4 --- /dev/null +++ b/lib/thread_sliced_shared.hpp @@ -0,0 +1,170 @@ +#ifndef __CHOREO_THREAD_SLICED_SHARED_HPP__ +#define __CHOREO_THREAD_SLICED_SHARED_HPP__ + +#include "ast.hpp" +#include "context.hpp" +#include "types.hpp" +#include "visitor.hpp" + +namespace Choreo { + +// Thread-sliced shared buffer analysis (GPU/CuTe target). +// +// A `shared` buffer whose content is thread-varying — i.e. any DMA moving data +// to or from the buffer indexes either side with a thread-level parallel +// variable (e.g. `dma.copy rhs.chunkat(k_tile, q#n_tile) => shared`) — cannot +// be a single block-common allocation: every thread stages its own slice. +// Without this analysis the CuTe codegen emitted one block-common buffer and +// loaded it under `__CHOREO_BLOCK_SINGLE__` (thread 0 only), so all threads +// computed from thread 0's slice — a silent miscompile. +// +// This pass records, for every such buffer, the number of per-thread slices +// (the bound of the enclosing thread-level parallel-by). Consumers: +// - MemAnalyzer (mem_reuse.cpp) scales the buffer footprint by the slice +// count so the shared spm holds the union of all slices; +// - MemReuse::ApplyMemOffset annotates the decl with the per-thread slice +// size ("tsliced" note) so codegen can add the per-thread addend +// `__choreo_vtid_x * ` to the buffer base pointer; +// - DMAPlan forces naive (per-thread) copies for DMAs touching the buffer +// and CuteCodeGen drops the `__CHOREO_BLOCK_SINGLE__` guard for them. +struct ThreadSlicedShared : public VisitorWithSymTab { +public: + ThreadSlicedShared() : VisitorWithSymTab("tsliced") {} + + // scoped buffer name -> number of per-thread slices (thread count). + static std::unordered_map& Store() { + static std::unordered_map store; + return store; + } + + static const ValueItem* Lookup(const std::string& scoped_name) { + auto& s = Store(); + auto it = s.find(scoped_name); + return it == s.end() ? nullptr : &it->second; + } + + static bool Enabled() { + return CCtx().MemReuse() && CCtx().TargetName() == "cute"; + } + +private: + std::unordered_map pv_levels_; + // bounds of enclosing thread-level parallel-bys (one 1-D level supported). + std::vector thread_bounds_; + +private: + bool BeforeVisitImpl(AST::Node& n) override { + if (isa(&n)) { + Store().clear(); + pv_levels_.clear(); + thread_bounds_.clear(); + } + if (auto pb = dyn_cast(&n)) { + pv_levels_[InScopeName(pb->BPV()->name)] = pb->GetLevel(); + for (auto id : pb->AllSubPVs()) + pv_levels_[InScopeName(cast(id)->name)] = + pb->GetLevel(); + if (pb->GetLevel() == ParallelLevel::THREAD) + thread_bounds_.push_back(pb->BoundValues()); + } + return true; + } + + bool AfterVisitImpl(AST::Node& n) override { + if (auto pb = dyn_cast(&n)) { + if (pb->GetLevel() == ParallelLevel::THREAD && !thread_bounds_.empty()) + thread_bounds_.pop_back(); + } + return true; + } + + // Does the chunkat reference any symbol deeper than block level? + // Sets `group_varying` when the reference is at (warp-)group level, which + // has its own (MMA) machinery and is out of scope for thread slicing. + bool IsThreadVarying(AST::ChunkAt* ca, bool& group_varying) const { + std::set syms; + if (ca->indices) { + auto s = ReferredSymbols(ca->indices.get(), this); + syms.insert(s.begin(), s.end()); + } + auto s = ReferredSymbols(ca, this); + syms.insert(s.begin(), s.end()); + syms.erase(InScopeName(ca->RefSymbol())); + bool varying = false; + for (auto& sym : syms) { + auto it = pv_levels_.find(sym); + if (it == pv_levels_.end()) continue; + if (it->second == ParallelLevel::THREAD) + varying = true; + else if (it->second == ParallelLevel::GROUP || + it->second == ParallelLevel::GROUPx4) + group_varying = true; + } + return varying; + } + + bool Visit(AST::DMA& n) override { + if (!Enabled()) return true; + if (n.operation == ".any") return true; + + auto fty = GetSpannedType(n.from->GetType()); + auto tty = GetSpannedType(n.to->GetType()); + if (!fty || !tty) return true; + bool from_shared = fty->GetStorage() == Storage::SHARED; + bool to_shared = tty->GetStorage() == Storage::SHARED; + if (!from_shared && !to_shared) return true; + + bool group_varying = false; + bool from_varying = false, to_varying = false; + if (auto ca = dyn_cast(n.from.get())) + from_varying = IsThreadVarying(ca, group_varying); + if (auto ca = dyn_cast(n.to.get())) + to_varying = IsThreadVarying(ca, group_varying); + bool varying = from_varying || to_varying; + if (varying && group_varying) + Error1(n.LOC(), + "thread-sliced shared buffers do not yet support mixing " + "thread-level and group-level parallel variables in DMA chunk " + "indexing."); + // Group-level-only variance uses the warp-level (MMA) machinery; leave + // the existing lowering untouched. + if (!varying || group_varying) return true; + + // Slicing is needed only when the DMA stages per-thread data INTO the + // shared buffer (thread-varying source: each thread's chunk is distinct + // data that must coexist), or writes per-thread results OUT to distinct + // global-memory regions (S2G with a thread-varying destination, which + // implies the buffer held per-thread data). A shared -> local copy with + // a thread-varying address is a distributed read of block-common data + // and needs no slicing. + bool slice_to = to_shared && from_varying; + bool slice_from = from_shared && to_varying && + (tty->GetStorage() == Storage::GLOBAL || + tty->GetStorage() == Storage::DEFAULT || to_shared); + if (!slice_to && !slice_from) return true; + if (thread_bounds_.size() != 1 || thread_bounds_.front().size() != 1) + Error1(n.LOC(), + "thread-sliced shared buffers do not yet support nested or " + "multi-dimensional thread-level parallel-bys."); + ValueItem count = thread_bounds_.front().front(); + if (!VIIsInt(count)) + Error1(n.LOC(), + "thread-varying DMA into shared storage requires an enclosing " + "thread-level parallel-by with a static bound."); + auto mark = [this, &count](AST::Node* side, bool shared_side) { + if (!shared_side) return; + auto ca = dyn_cast(side); + if (!ca) return; + auto sname = InScopeName(ca->RefSymbol()); + auto& s = Store(); + if (!s.count(sname)) s.emplace(sname, count); + }; + mark(n.from.get(), slice_from); + mark(n.to.get(), slice_to); + return true; + } +}; + +} // namespace Choreo + +#endif // __CHOREO_THREAD_SLICED_SHARED_HPP__ diff --git a/lib/valno.cpp b/lib/valno.cpp index 31fd5ac8..415273c4 100644 --- a/lib/valno.cpp +++ b/lib/valno.cpp @@ -1,4 +1,5 @@ #include "shapeinfer.hpp" +#include "context.hpp" using namespace Choreo; using namespace Choreo::valno; @@ -352,6 +353,11 @@ const SignTy ValueNumbering::TryToSimplifyBinary(const OpTy& op, auto l_csn = CSign(lhs); auto r_csn = CSign(rhs); + // Ablation (--disable-vn-simplify): skip the simplification rules that + // rely on value-number equality and bind sets; constant folding and the + // structural sbe normalization above stay enabled. + const bool vn_simp = !CCtx().DisableVNSimplify(); + auto Report = [this, &op, &lhs, &rhs, &verbose](const SignTy& s) -> const SignTy { if (trace && verbose) @@ -372,7 +378,7 @@ const SignTy ValueNumbering::TryToSimplifyBinary(const OpTy& op, // or else, apply optimization for the symbolic expression else if (op == Op::Div) { // a/a == 1 - if (NumSign(lhs) == NumSign(rhs)) { + if (vn_simp && NumSign(lhs) == NumSign(rhs)) { return Report(c_sn(1)); } // a/1 = a @@ -380,7 +386,7 @@ const SignTy ValueNumbering::TryToSimplifyBinary(const OpTy& op, rc && rc->Holds() && rc->GetInt() == 1) return Report(lhs); // useful simplification: a/(a/b) = b - else if (lhs->Count() == 1 /*not multiple values*/) { + else if (vn_simp && lhs->Count() == 1 /*not multiple values*/) { NumTy rvn = GetValueNumberOfSignature(rhs); auto bind_set = GetBindSet(rvn); bind_set.insert(rvn); // always add self @@ -394,10 +400,10 @@ const SignTy ValueNumbering::TryToSimplifyBinary(const OpTy& op, } } else if (op == Op::Sub) { // a-a == 0 - if (NumSign(lhs) == NumSign(rhs)) return Report(c_sn(0)); + if (vn_simp && NumSign(lhs) == NumSign(rhs)) return Report(c_sn(0)); } else if (op == Op::Add) { // useful simplification: a-b+b = a - if (rhs->Count() == 1 /*not multiple values*/) { + if (vn_simp && rhs->Count() == 1 /*not multiple values*/) { NumTy lvn = NumSign(lhs); auto bind_set = GetBindSet(lvn); bind_set.insert(lvn); // always add self @@ -415,7 +421,7 @@ const SignTy ValueNumbering::TryToSimplifyBinary(const OpTy& op, // if `xx.chunkat(a#b)`, then the result shape should be 1 // that is, N / (#a * #b) = N / N = 1 // so, `a#b` should be simplified to a bounded var whose ubound is `N` - if (lhs->Count() == 1 /*not multiple values*/) { + if (vn_simp && lhs->Count() == 1 /*not multiple values*/) { NumTy rvn = NumSign(rhs); auto bind_set = GetBindSet(rvn); bind_set.insert(rvn); // always add self @@ -426,7 +432,7 @@ const SignTy ValueNumbering::TryToSimplifyBinary(const OpTy& op, assert(div.size() == 2); if (NumSign(lhs) == NumSign(div[1])) return Report(div[0]); } - } else if (rhs->Count() == 1) { + } else if (vn_simp && rhs->Count() == 1) { // TODO: # is different with * // a # (b/a) will alway result in a? // if so, we need to emphasize this optimization to our users. @@ -480,6 +486,15 @@ ValueNumbering::GetOrGenValueNumberFromSignature(const SignTy& signature) { if (IsUnknown(signature)) return NumTy::Unknown(); if (IsNone(signature)) return NumTy::None(); + // Ablation (--disable-vn-share): generate a fresh, unshared valno for + // every operation-expression occurrence. Constants, symbols and shape + // tuples keep their shared identity so that signature re-derivation stays + // consistent. GenerateFresh keeps a first-occurrence record in sign_pool + // so that NumSign() lookups remain total; bypassing FindValueNum here is + // what disables cross-occurrence sharing. + if (CCtx().DisableVNShare() && isa(signature)) + return vntbl.GenerateFresh(signature); + if (auto found = vntbl.FindValueNum(signature)) return *found; return GenerateValueNumberFromSignature(signature); } diff --git a/lib/valno.hpp b/lib/valno.hpp index 693a395c..316162a8 100644 --- a/lib/valno.hpp +++ b/lib/valno.hpp @@ -806,6 +806,22 @@ class ValueNumberTable { return valno; } + // Generate a fresh valno without interning for sharing. + // Used by --disable-vn-share ablation mode. The first fresh valno is + // recorded in sign_pool (emplace is a no-op if an earlier occurrence + // already recorded one) so that NumSign()/GetValueNumberOfSignature stay + // total and stable; sharing is still disabled because the ablation hook + // in GetOrGenValueNumberFromSignature bypasses the pools entirely for + // operation signatures. The signature is also kept in value_nums so that + // NumSign()/SignNum() can still retrieve it. + NumTy GenerateFresh(const SignTy& s) { + auto valno = next_valno++; + assert(value_nums.count(valno) == 0); + value_nums.emplace(valno, std::vector{s}); + sign_pool.emplace(s, valno); + return valno; + } + // Bind a valno to the existing (dummy) signature void BindDummy(const SignTy& s, NumTy v) { // note: only dummy sign can be re-generated diff --git a/runtime/choreo.h b/runtime/choreo.h index 279977e6..11d61b24 100644 --- a/runtime/choreo.h +++ b/runtime/choreo.h @@ -37,8 +37,14 @@ #include "private_target0_defines.h" #endif +// Fallback abort when the target does not provide its own __co_abort__. +// __builtin_trap is host-only under nvcc; device code must use __trap(). #ifndef __co_abort__ - #define __co_abort__() __builtin_trap() + #if defined(__CUDA_ARCH__) + #define __co_abort__() __trap() + #else + #define __co_abort__() __builtin_trap() + #endif #endif #ifdef __CHOREO_PRIVATE_TGT0__ @@ -394,8 +400,7 @@ class ArrayProxy { template typename std::enable_if<(M == 1), T&>::type // make sure to return the reference type - __co_any__ - operator[](int index) { + __co_any__ operator[](int index) { choreo_assert(index >= 0, "Index out of bounds", __FILE__, __LINE__); choreo_assert((size_t)index < (*dims)[0], "Index out of bounds", __FILE__, __LINE__); @@ -782,8 +787,7 @@ class spanned_view { template typename std::enable_if<(M == 1), T&>::type // make sure to return the reference type - __co_any__ - operator[](int index) { + __co_any__ operator[](int index) { choreo_assert(index >= 0, "Index out of bounds", __FILE__, __LINE__); choreo_assert((size_t)index < dims[0], "Index out of bounds", __FILE__, __LINE__); diff --git a/runtime/choreo_cute.h b/runtime/choreo_cute.h index c1577153..4ead9a88 100644 --- a/runtime/choreo_cute.h +++ b/runtime/choreo_cute.h @@ -517,9 +517,15 @@ struct TMAAtom { using AsyncCopyAtom = cute::AutoCopyAsync; + // choreo.h defines __co_abort__ as a macro (falling back to __builtin_trap) + // when the target does not provide its own; defining the function below in + // that case would macro-expand to a choreo::__builtin_trap declaration and + // make every abort call ambiguous under nvcc. + #ifndef __co_abort__ __device__ __attribute__((always_inline)) static inline void __co_abort__() { __trap(); } + #endif // this facilitate the wait-N implementation for sm_80+ // it is must be warp-wise since async copy is warp-wise. diff --git a/tests/gpu/codegen/cute/thread_sliced_shared.co b/tests/gpu/codegen/cute/thread_sliced_shared.co new file mode 100644 index 00000000..bb51f0cf --- /dev/null +++ b/tests/gpu/codegen/cute/thread_sliced_shared.co @@ -0,0 +1,38 @@ +// RUN: choreo -es -t cute -arch=sm_86 %s -o - | FileCheck %s + +// Thread-sliced shared buffers: a DMA whose chunk indexing depends on a +// thread-level parallel variable stages one slice per thread in shared +// memory (union-of-slices). The buffer footprint is scaled by the thread +// count, the base pointer carries a per-thread addend, and every thread +// performs its own copy -- a __CHOREO_BLOCK_SINGLE__ guard would move only +// thread 0's slice (the previous, silently-wrong lowering). + +__co__ auto foo(f32 [64, 64] input) { + f32 [input.span(0), input.span(1)] output; + parallel p by 4, q by 8 { + shared f32[input.span(0)/#p/#q, input.span(1)] l1_out {0.0f}; + l1_a = dma.copy input.chunkat(p#q, _) => shared; + foreach {i} in [l1_out.span(0)] { + foreach {j} in [l1_out.span(1)] { + l1_out.at(i, j) = l1_a.data.at(i, j) * 2.0f; + } + } + dma.copy l1_out => output.chunkat(p#q, _); + } + return output; +} + +// CHECK: __global__ void __choreo_device_foo(float * input, float * output) { +// footprint: two [2, 64] f32 buffers (512 B per thread slice) x 8 threads. +// CHECK: __shared__ alignas(16) unsigned char anon_{{[0-9]+}}[8192]; +// CHECK-NOT: __CHOREO_BLOCK_SINGLE__ +// CHECK: float* l1_out = (float*)(anon_{{[0-9]+}} + 4096 + __choreo_vtid_x * (512)); +// per-thread slice initialization (unguarded). +// CHECK: for (int i = 0; i < 128; ++i) l1_out[i] = 0.000000f; +// CHECK-NOT: __CHOREO_BLOCK_SINGLE__ +// CHECK: float* l1_a__buf__ = (float*)(anon_{{[0-9]+}} + 0 + __choreo_vtid_x * (512)); +// per-thread G2S and S2G copies (unguarded). +// CHECK-NOT: __CHOREO_BLOCK_SINGLE__ +// CHECK: choreo::naive_copy(__tensor1_input, __tensor2_l1_a__buf__); +// CHECK-NOT: __CHOREO_BLOCK_SINGLE__ +// CHECK: choreo::naive_copy(__tensor3_l1_out, __tensor4_output);