[Ascend] Support partial L0 copies with explicit stride

src/ascend/layout/ascend_layouts.cc

@@ -5,12 +5,14 @@
  */
 
 #include "support/check.h"
+#include <algorithm>
 #include <tvm/arith/analyzer.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/runtime/logging.h>
 #include <tvm/tirx/op.h>
 
 #include "ascend/layout/ascend_layouts.h"
+#include "ascend/op/utils.h"
 #include "layout/layout.h"
 
 namespace tvm {
@@ -380,6 +382,28 @@ bool TryExtractAscendFractalRegion(const Layout &layout,
   return true;
 }
 
+BufferRegion ExpandAscendFractalMayAccess(const BufferRegion &region) {
+  const Buffer &buffer = region->buffer;
+  int bits = buffer->dtype.bits();
+  if (!(IsL1Buffer(buffer) || IsL0ABuffer(buffer) || IsL0BBuffer(buffer) ||
+        IsL0CBuffer(buffer)) ||
+      region->region.size() < 2 || bits <= 0 || 256 % bits != 0)
+    return region;
+  int alignment = IsL0CBuffer(buffer) ? 16 : std::max(16, 256 / bits);
+  ffi::Array<Range> ranges = region->region;
+  for (const Range &range : ranges)
+    if (is_zero(range->extent))
+      return region;
+  for (size_t i = ranges.size() - 2; i < ranges.size(); ++i) {
+    PrimExpr begin = floordiv(ranges[i]->min, alignment) * alignment;
+    PrimExpr end = floordiv(ranges[i]->min + ranges[i]->extent + alignment - 1,
+                            alignment) *
+                   alignment;
+    ranges.Set(i, Range::FromMinExtent(begin, end - begin));
+  }
+  return BufferRegion(buffer, ranges);
+}
+
 // Ascend ND (row-major) identity layout for GM/UB operands. Forward is the
 // identity [i0, .., i_{n-1}] -> [i0, .., i_{n-1}], so OutputShape ==
 // InputShape. Used only for copy-parameter derivation in lowering;

src/ascend/layout/ascend_layouts.h

@@ -38,7 +38,7 @@
 
 // Ascend ND (row-major) identity layout for GM/UB operands; OutputShape ==
 // InputShape. Used for copy-parameter derivation only (not anchored).
 Layout makeAscendNDLayout(const tirx::Buffer &buffer);
 
 /// Compute a 0-origin strided buffer view for a logical sub-region on a
 /// layout-annotated physical buffer.
@@ -127,6 +127,24 @@
                                    const ffi::Array<Range> &logical_region,
                                    AscendFractalRegionInfo *out);
 
+/// Conservatively expand an L1/L0 data region in logical coordinates to cover
+/// whole fractals for either Cube major. Both trailing axes are rounded to
+/// max(16, C0), or 16 for L0C; leading axes and statically empty regions remain
+/// unchanged. Unsupported scopes, ranks, or element widths pass through.
+/// This is a may-access bound only: it must not establish must-write coverage
+/// or determine the actual instruction transfer geometry.
+tirx::BufferRegion
+ExpandAscendFractalMayAccess(const tirx::BufferRegion &region);
+
+namespace attr {
+
+// Copy lowering transfers the destination pitch to its tvm_access_ptr.
+// LowerTileOp uses it to remap matrix offsets while retaining the allocation's
+// leading-plane / buffer-version stride.
+constexpr const char *kL0AccessPitch = "tl.ascend_l0_access_pitch";
+
+} // namespace attr
+
 } // namespace tl
 } // namespace tvm
 

src/ascend/op/copy.h

@@ -39,7 +39,8 @@ class AscendCopyNode : public CopyNode {
   PrimExpr unit_flag_ctl;          // "unit_flag_ctrl", defaults to 0
   PrimExpr sub_blockid;            // "sub_blockid", defaults to 0
   Optional<Integer> l2_cache_ctrl; // "l2_cache_ctrl"; defaults differ per path
-  Optional<PrimExpr> pad_value;    // "pad_value": GM->UB pad fill value
+  Optional<PrimExpr> stride;    // "stride": L0 matrix pitch in fractal blocks
+  Optional<PrimExpr> pad_value; // "pad_value": GM->UB pad fill value
 
   /*!
    * \brief l2_cache_ctrl with a path-specific default (GM loads default to 0,

src/ascend/transform/auto_schedule/dependency_analysis.cc

@@ -38,6 +38,7 @@
 #include <tvm/tirx/expr.h>
 #include <tvm/tirx/stmt_functor.h>
 
+#include "ascend/layout/ascend_layouts.h"
 #include "support/check.h"
 #include "transform/common/attr.h"
 #include "transform/common/constr_visitor.h"
@@ -116,7 +117,9 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region,
                         const ConstrSet &b_ctx, const BufferRegion &b_region,
                         ControlNode *loop, int offset,
                         size_t num_storage_owners,
-                        const ConflictHintList &root_conflicts) {
+                        const ConflictHintList &root_conflicts,
+                        Optional<PrimExpr> a_pitch,
+                        Optional<PrimExpr> b_pitch) {
   const bool cross = offset != 0;
   if (cross && loop == nullptr)
     return false;
@@ -150,8 +153,8 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region,
   if (!a_region->buffer.same_as(b_region->buffer))
     return true;
 
-  const Array<Range> &a_rng = a_region->region;
-  const Array<Range> &b_rng = b_region->region;
+  Array<Range> a_rng = ExpandAscendFractalMayAccess(a_region)->region;
+  Array<Range> b_rng = ExpandAscendFractalMayAccess(b_region)->region;
   if (a_rng.size() != b_rng.size())
     return true; // shape mismatch — be safe
   const size_t ndim = a_rng.size();
@@ -208,21 +211,35 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region,
   }
   loop_ctx.Merge(a_local_ctx).Merge(b_local_ctx).Populate(ana);
 
-  // Freshen mutable reads (BufferLoads etc.) inline in the region bounds, per
-  // side: a read written verbatim in both bounds is not provably equal across
-  // the two access points (a store may sit between them), so it must become an
-  // independent unknown on each side.
   FreshenMutableReads freshen_a(FreshenMutableReads::Mode::kSnapshot);
   FreshenMutableReads freshen_b(FreshenMutableReads::Mode::kSnapshot);
+  size_t compared_dims = ndim;
+  const Buffer &buffer = a_region->buffer;
+  if (IsL0ABuffer(buffer) || IsL0BBuffer(buffer) || IsL0CBuffer(buffer)) {
+    bool same_pitch = false;
+    if (a_pitch.defined() && b_pitch.defined()) {
+      PrimExpr equal = freshen_a(Substitute(a_pitch.value(), a_sub)) ==
+                       freshen_b(Substitute(b_pitch.value(), b_sub));
+      same_pitch = ana.CanProve(equal) || ana.z3_prover.CanProve(equal);
+    }
+    // L0 matrix regions may overlap unless their pitches are provably equal.
+    if (!same_pitch)
+      compared_dims = ndim > 2 ? ndim - 2 : 0;
+  }
   PrimExpr any_dim_disjoint = Bool(false);
-  for (size_t i = 0; i < ndim; ++i) {
+  for (size_t i = 0; i < compared_dims; ++i) {
     PrimExpr a_min = freshen_a(Substitute(a_rng[i]->min, a_sub));
-    PrimExpr a_max =
-        freshen_a(Substitute(a_rng[i]->min + a_rng[i]->extent - 1, a_sub));
+    PrimExpr a_extent = freshen_a(Substitute(a_rng[i]->extent, a_sub));
     PrimExpr b_min = freshen_b(Substitute(b_rng[i]->min, b_sub));
-    PrimExpr b_max =
-        freshen_b(Substitute(b_rng[i]->min + b_rng[i]->extent - 1, b_sub));
-    any_dim_disjoint = any_dim_disjoint || (a_max < b_min) || (b_max < a_min);
+    PrimExpr b_extent = freshen_b(Substitute(b_rng[i]->extent, b_sub));
+    // A half-open endpoint can exceed the int32 index range by one even
+    // for an ordinary single-element access.
+    PrimExpr a_end =
+        cast(DataType::Int(64), a_min) + cast(DataType::Int(64), a_extent);
+    PrimExpr b_end =
+        cast(DataType::Int(64), b_min) + cast(DataType::Int(64), b_extent);
+    any_dim_disjoint = any_dim_disjoint || (a_extent <= 0) || (b_extent <= 0) ||
+                       (a_end <= b_min) || (b_end <= a_min);
   }
   any_dim_disjoint = ana.Simplify(any_dim_disjoint);
   // Raw Z3 as a fallback catches the infeasible-context case that Analyzer's
@@ -389,8 +406,9 @@ AnalyzeDependencies(std::vector<IRStructure *> nodes, ControlNode *loop,
               for (int d = (i < j ? 0 : 1); d <= manual_versions; ++d) {
                 if (RegionsMayConflict(lhs.task->outer_ctx, lhs.region,
                                        rhs.task->outer_ctx, rhs.region, loop, d,
-                                       /*num_storage_owners=*/0,
-                                       root_conflicts)) {
+                                       /*num_storage_owners=*/0, root_conflicts,
+                                       lhs.task->GetL0Pitch(lhs_storage),
+                                       rhs.task->GetL0Pitch(rhs_storage))) {
                   add_pair(storage, d, lhs.task, rhs.task);
                   break;
                 }
@@ -405,7 +423,9 @@ AnalyzeDependencies(std::vector<IRStructure *> nodes, ControlNode *loop,
               if (RegionsMayConflict(lhs.task->outer_ctx, lhs.region,
                                      rhs.task->outer_ctx, rhs.region, loop,
                                      distance, get_num_storage_owners(storage),
-                                     root_conflicts)) {
+                                     root_conflicts,
+                                     lhs.task->GetL0Pitch(lhs_storage),
+                                     rhs.task->GetL0Pitch(rhs_storage))) {
                 add_pair(storage, owner_exclusion ? 0 : distance, lhs.task,
                          rhs.task,
                          owner_exclusion ? DependencyKind::kOwnerExclusion

src/ascend/transform/auto_schedule/dependency_analysis.h

@@ -109,7 +109,8 @@ AnalyzeDependencies(std::vector<IRStructure *> nodes,
                     const ConflictHintList &root_conflicts = {});
 
 // Check whether two task regions may overlap under their respective symbolic
-// contexts. `offset == 0` compares accesses in the same iteration. A positive
+// contexts. Only a proof of disjointness removes a conflict.
+// `offset == 0` compares accesses in the same iteration. A positive
 // offset constrains the current consumer loop coordinate to equal the producer
 // coordinate plus that offset. A negative offset denotes an automatic
 // cross-iteration check. Normally it renames the loop suffix below the nearest
@@ -123,7 +124,9 @@ bool RegionsMayConflict(const ConstrSet &a_ctx, const BufferRegion &a_region,
                         const ConstrSet &b_ctx, const BufferRegion &b_region,
                         ControlNode *loop, int offset,
                         size_t num_storage_owners = 0,
-                        const ConflictHintList &root_conflicts = {});
+                        const ConflictHintList &root_conflicts = {},
+                        Optional<PrimExpr> a_pitch = std::nullopt,
+                        Optional<PrimExpr> b_pitch = std::nullopt);
 
 } // namespace ascend
 } // namespace tl

src/ascend/transform/auto_schedule/ir_structure.cc

@@ -299,6 +299,7 @@ std::shared_ptr<IRStructure> TaskNode::Clone() const {
   // Copy memory access regions
   new_task->SetReadRegions(GetReadRegions());
   new_task->SetWriteRegions(GetWriteRegions());
+  new_task->SetL0Pitches(GetL0Pitches());
   new_task->SetReadVars(GetReadVars());
   new_task->SetWriteVars(GetWriteVars());
   new_task->outer_ctx = outer_ctx;

src/ascend/transform/auto_schedule/ir_structure.h

@@ -556,6 +556,16 @@ class TaskNode : public IRStructure {
   void SetWriteRegions(const std::vector<BufferRegion> &regions) {
     write_regions_ = regions;
   }
+  const Map<Var, Optional<PrimExpr>> &GetL0Pitches() const {
+    return l0_pitches_;
+  }
+  Optional<PrimExpr> GetL0Pitch(const Var &storage) const {
+    auto it = l0_pitches_.find(storage);
+    return it == l0_pitches_.end() ? Optional<PrimExpr>() : (*it).second;
+  }
+  void SetL0Pitches(const Map<Var, Optional<PrimExpr>> &pitches) {
+    l0_pitches_ = pitches;
+  }
   void SetReadVars(const std::vector<Var> &vars) { read_vars_ = vars; }
   void SetWriteVars(const std::vector<Var> &vars) { write_vars_ = vars; }
   void SubstituteVar(const Var &old_var, const Var &new_var) override {
@@ -588,6 +598,14 @@ class TaskNode : public IRStructure {
     substitute_vars(write_vars_);
     substitute_regions(read_regions_);
     substitute_regions(write_regions_);
+    Map<Var, Optional<PrimExpr>> pitches;
+    for (const auto &[storage, pitch] : l0_pitches_) {
+      pitches.Set(storage.same_as(old_var) ? new_var : storage,
+                  pitch.defined() ? Optional<PrimExpr>(
+                                        Substitute(pitch.value(), substitution))
+                                  : std::nullopt);
+    }
+    l0_pitches_ = std::move(pitches);
     outer_ctx = outer_ctx.Substitute(substitution);
     stmt = Substitute(stmt, substitution);
     SubstituteVarInGuards(old_var, new_var);
@@ -657,6 +675,7 @@ class TaskNode : public IRStructure {
   // Logical access regions. Physical footprints are derived from current IR.
   std::vector<BufferRegion> read_regions_;
   std::vector<BufferRegion> write_regions_;
+  Map<Var, Optional<PrimExpr>> l0_pitches_;
   std::vector<Var> read_vars_;
   std::vector<Var> write_vars_;