[Ascend] Preserve buffer view layouts on shared storage

.agents/skills/tilelang-ascend/SKILL.md

@@ -347,6 +347,21 @@ when changing manual kernels. For capacity or launch failures, record the
 generated allocation sizes, buffer versions, VF mode, and compiler options.
 Check SDK reservations as well as user buffers before changing tile sizes.
 
+For Cube storage (L1, L0A, L0B, L0C), `T.view` and `T.reshape` aliases can have
+different canonical layouts inferred from their operations or supplied with
+`T.annotate_layout`. For example, a `(64, 128)` L0B allocation can have a
+`(128, 64)` view filled with `T.copy(..., transpose=True)` and consumed with
+`T.gemm(..., transpose_B=True)`. A same-shape view can instead use a different
+`MajorK`/`MajorMN` layout. Each view keeps its own address mapping while the
+shared `.data` keeps one allocation and version clock; `T.assume_conflict` is
+unnecessary for these aliases.
+
+These views require dense logical strides, zero base offset, the same accessed
+dtype, and equal padded physical sizes. They reinterpret storage without moving
+existing values: the producer must write the representation its consumer uses.
+Use `T.access_ptr(view[...], ...)` when passing a view to an intrinsic, so layout
+lowering can select its address mapping.
+
 `T.annotate_unlimit_memory("shared", "shared.l1")` inside a kernel removes the
 scheduler's capacity constraint for those pools. Other supported scopes are
 `shared.l0a`, `shared.l0b`, and `shared.l0c`. This is useful for investigating a

src/ascend/op/copy.cc

@@ -423,11 +423,13 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T,
   Array<Range> src_range = NormalizeEmptyUnitAxesForMTE(op.src_range, analyzer);
   Array<Range> dst_range = NormalizeEmptyUnitAxesForMTE(op.dst_range, analyzer);
 
-  auto make_access_ptr = [&](const Buffer &buf, PrimExpr offset,
+  auto make_access_ptr = [&](const Buffer &buf, const Array<Range> &ranges,
                              PrimExpr extent, int rw_mask) -> PrimExpr {
-    return Call(
-        DataType::Handle(), builtin::tvm_access_ptr(),
-        {TypeAnnotation(buf->dtype), buf->data, offset, extent, I(rw_mask)});
+    Array<PrimExpr> indices;
+    for (const Range &range : ranges)
+      indices.push_back(range->min);
+    return Call(DataType::Handle(), tl::access_ptr(),
+                {BufferLoad(buf, indices), extent, I(rw_mask)});
   };
   auto region_elements = [&](const Array<Range> &ranges) {
     PrimExpr result = I(1);
@@ -442,10 +444,8 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T,
       op.dst, dst_range, op.dst->dtype.bits() * op.dst->dtype.lanes());
   PrimExpr src_region_elements = region_elements(src_range);
   PrimExpr dst_region_elements = region_elements(dst_range);
-  PrimExpr src_ptr =
-      make_access_ptr(op.src, src_region.offset, src_region_elements, 1);
-  PrimExpr dst_ptr =
-      make_access_ptr(op.dst, dst_region.offset, dst_region_elements, 2);
+  PrimExpr src_ptr = make_access_ptr(op.src, src_range, src_region_elements, 1);
+  PrimExpr dst_ptr = make_access_ptr(op.dst, dst_range, dst_region_elements, 2);
 
   PrimExpr sid = I(0);
   PrimExpr zero = I(0);
@@ -464,9 +464,9 @@ Stmt LowerDMACopy(const AscendCopyNode &op, const LowerArgs &T,
 
     PrimExpr plan_total_elems = plan.row_elems * plan.n_rows;
     PrimExpr plan_src_ptr =
-        make_access_ptr(op.src, plan.src_offset, plan_total_elems, 1);
+        make_access_ptr(op.src, src_range, plan_total_elems, 1);
     PrimExpr plan_dst_ptr =
-        make_access_ptr(op.dst, plan.dst_offset, plan_total_elems, 2);
+        make_access_ptr(op.dst, dst_range, plan_total_elems, 2);
     PrimExpr plan_total_bytes = plan_row_bytes * plan.n_rows;
 
     // Ascend GM->UB padding (data_select): an unaligned row is right-padded up

src/ascend/op/oob_padding.cc

@@ -69,15 +69,7 @@ DMAPath GetDMAPath(const Buffer &src, const Buffer &dst) {
 
 Optional<Layout> FindLayoutForBuffer(const LayoutMap &layout_map,
                                      const Buffer &buffer) {
-  if (layout_map.count(buffer)) {
-    return layout_map[buffer];
-  }
-  for (const auto &kv : layout_map) {
-    if (kv.first->data.same_as(buffer->data)) {
-      return kv.second;
-    }
-  }
-  return std::nullopt;
+  return layout_map.Get(buffer);
 }
 
 BoundedDMACopyRanges ClampDMACopyTail(const AscendCopyNode &op,

src/ascend/op/oob_padding.h

@@ -26,7 +26,7 @@
 namespace tl {
 namespace ascend {
 
 using namespace tirx;
 
 /*! \brief Ascend DMA route selected from the source/destination scopes. */
 enum class DMAPath {
@@ -53,8 +53,8 @@
 DMAPath GetDMAPath(const Buffer &src, const Buffer &dst);
 
 /*!
- * \brief Look up a buffer's inferred Layout in a layout map, matching either
- * the buffer object or its underlying data Var. Returns nullopt if absent.
+ * \brief Look up a view's inferred Layout by Buffer identity. Storage aliases
+ * may have distinct layouts; inference resolves any layout inheritance.
  */
 Optional<Layout> FindLayoutForBuffer(const LayoutMap &layout_map,
                                      const Buffer &buffer);

src/ascend/op/utils.h

@@ -13,6 +13,7 @@
 
 #include <algorithm>
 
+#include "op/builtin.h"
 #include "op/utils.h"
 
 namespace tvm {
@@ -75,7 +76,7 @@ inline bool IsAscendOnChipBuffer(const Buffer &buffer) {
 }
 
 /*!
- * \brief tvm_access_ptr for a region whose innermost 2-D position travels as
+ * \brief A view-preserving pointer whose innermost 2-D position travels as
  *        separate intrinsic arguments.
  *
  * The Ascend L1->L0 loads pass the M/K start positions as their own arguments
@@ -85,13 +86,8 @@ inline bool IsAscendOnChipBuffer(const Buffer &buffer) {
  * example -- and counting the last two dims into the offset as well would place
  * that position twice.
  *
- * This builds exactly that pointer, on top of the neutral
- * MakeAccessPtrFromRegion: the innermost two dims are re-based to zero before
- * delegating, which leaves their extents (and hence the reported extent) alone
- * and drops them from the offset.
- *
- * The buffer keeps ``require_2d == false`` semantics, so 1-D regions still take
- * the neutral function's dedicated path and are unaffected.
+ * Keep the Buffer identity until layout lowering: a data Var alone cannot
+ * identify which canonical matrix view supplies the physical addresses.
  */
 inline PrimExpr MakeAscendLeadingDimAccessPtr(const BufferRegion &region,
                                               int rw_mask) {
@@ -104,8 +100,15 @@ inline PrimExpr MakeAscendLeadingDimAccessPtr(const BufferRegion &region,
     const auto &r = leading[i];
     leading.Set(i, Range::FromMinExtent(make_zero(r->min.dtype()), r->extent));
   }
-  return MakeAccessPtrFromRegion(BufferRegion(region->buffer, leading),
-                                 rw_mask);
+  ffi::Array<PrimExpr> indices;
+  for (const Range &range : leading)
+    indices.push_back(range->min);
+  PrimExpr extent = leading.back()->extent;
+  if (n >= 2)
+    extent = extent * leading[n - 2]->extent;
+  return Call(DataType::Handle(), tl::access_ptr(),
+              {BufferLoad(region->buffer, indices), extent,
+               IntImm(DataType::Int(32), rw_mask)});
 }
 
 } // namespace tl

src/ascend/transform/layout_inference.cc

@@ -316,6 +316,11 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer {
       // Helper: propagate inferred layout to alias buffers (same data Var)
       auto propagate_alias = [&](const Buffer &src_buffer,
                                  const Layout &src_layout) {
+        // Fragment aliases participate in the thread-ownership solver.
+        // Memory layouts belong to individual views; fill missing layouts
+        // only after each view's own constraints have been inferred.
+        if (!src_layout.as<FragmentNode>())
+          return;
         if (!buffer_data_to_buffers_.count(src_buffer->data))
           return;
         const auto &siblings = buffer_data_to_buffers_[src_buffer->data];
@@ -623,45 +628,41 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer {
                      in_queue);
     // step 3: relax constraints to free and re-run
     InferInFreeMode(layout_map, strict_layout_map);
-    // step 4: finalize alias layouts by Var
-    // For each storage var, if any buffer in the group has a layout,
-    // propagate (reshape if needed) to the rest to ensure completeness.
+    // step 4: fill unconstrained views only when all known interpretations of
+    // their storage agree. A shared data Var does not impose layout equality.
+    const LayoutMap inferred_layouts = layout_map;
     for (const auto &[var, buffers] : buffer_data_to_buffers_) {
-      // Find a representative with existing layout
-      Optional<Buffer> rep;
-      Optional<Layout> rep_layout;
-      for (const auto &buf : buffers) {
-        if (layout_map.count(buf)) {
-          rep = buf;
-          rep_layout = layout_map[buf];
-          break;
-        }
-      }
-      if (!rep_layout.defined())
-        continue;
       for (const auto &buf : buffers) {
-        if (!layout_map.count(buf)) {
-          bool shapes_equal =
-              rep_layout.value()->InputShape().size() == buf->shape.size();
-          if (shapes_equal) {
-            for (size_t i = 0; i < rep_layout.value()->InputShape().size();
-                 ++i) {
-              if (!analyzer_.CanProveEqual(rep_layout.value()->InputShape()[i],
-                                           buf->shape[i])) {
-                shapes_equal = false;
-                break;
-              }
-            }
-          }
-          Layout reshaped =
-              shapes_equal
-                  ? rep_layout.value()
-                  : rep_layout.value()->Reshape(
+        if (layout_map.count(buf))
+          continue;
+        Optional<Layout> inherited;
+        bool ambiguous = false;
+        for (const auto &source : buffers) {
+          Optional<Layout> known = inferred_layouts.Get(source);
+          if (!known.defined())
+            continue;
+          Layout candidate =
+              ShapesEqual(known.value()->InputShape(), buf->shape, &analyzer_)
+                  ? known.value()
+                  : known.value()->Reshape(
                         buf->shape, &analyzer_,
-                        Integer(GetElementStorageBits(rep.value()->dtype)),
+                        Integer(GetElementStorageBits(source->dtype)),
                         Integer(GetElementStorageBits(buf->dtype)));
-          layout_map.Set(buf, reshaped);
+          if (inherited.defined() &&
+              !inherited.value()->IsEqual(candidate.get())) {
+            ambiguous = true;
+            break;
+          }
+          inherited = candidate;
         }
+        ICHECK(!ambiguous || !accessed_buffers_.count(buf))
+            << "Ambiguous layout for buffer view " << buf->name << ": storage "
+            << var->name_hint
+            << " has distinct layouts. Annotate this view explicitly.";
+        // An unused backing Buffer only supplies allocation geometry. Any
+        // candidate suffices; storage normalization checks physical footprints.
+        if (inherited.defined())
+          layout_map.Set(buf, inherited.value());
       }
     }
 
@@ -869,6 +870,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer {
   }
 
   void addToUseList(const Buffer &buffer) {
+    accessed_buffers_.insert(buffer);
     // Fragment buffers and reducers (whose PartialFragment travels the same
     // use_list_ edges: commits notify users, shared buffers union free-mode
     // components) both register; anything else has no layout to solve.
@@ -1001,18 +1003,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer {
       // annotated layouts are available during InferLayout.
       // Must happen after visiting body so buffer_data_to_buffers_ is
       // populated.
-      if (pushed_region_scope && op->annotations.count(attr::kLayoutMap)) {
-        auto map = op->annotations.Get(attr::kLayoutMap)
-                       ->as<Map<Var, Layout>>()
-                       .value();
-        for (const auto &[var, layout] : map) {
-          if (buffer_data_to_buffers_.count(var)) {
-            for (const auto &buffer : buffer_data_to_buffers_[var]) {
-              annotated_layout_map_.Set(buffer, layout);
-            }
-          }
-        }
-      }
+      RecordLayoutAnnotations(GetRef<SBlock>(op));
       // Regions contribute nothing to the inference worklist themselves:
       // nested tile operators were already collected while visiting the
       // body, and the region block is consumed by codegen as-is.
@@ -1035,16 +1026,23 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer {
     // First, visit the block body to collect all buffers from
     // BufferLoad/BufferStore
     IRVisitorWithAnalyzer::VisitStmt_(op);
+    RecordLayoutAnnotations(GetRef<SBlock>(op));
+  }
 
+  void RecordLayoutAnnotations(const SBlock &block) {
     // After visiting, apply layouts to all collected buffers
-    if (op->annotations.count(attr::kLayoutMap)) {
-      // Check if the layout map is Map<Var, Layout>
-      auto map =
-          op->annotations.Get(attr::kLayoutMap)->as<Map<Var, Layout>>().value();
-      for (const auto &[var, layout] : map) {
+    if (block->annotations.count(attr::kLayoutMap)) {
+      auto map = block->annotations.Get(attr::kLayoutMap)
+                     ->cast<Map<ObjectRef, Layout>>();
+      for (const auto &[key, layout] : map) {
+        Optional<Buffer> annotated = key.as<Buffer>();
+        Var var =
+            annotated.defined() ? annotated.value()->data : Downcast<Var>(key);
         ICHECK(buffer_data_to_buffers_.count(var))
             << "buffer " << var << " is not found in the block";
-        const auto &buffers = buffer_data_to_buffers_[var];
+        Array<Buffer> buffers = annotated.defined()
+                                    ? Array<Buffer>{annotated.value()}
+                                    : buffer_data_to_buffers_[var];
         ICHECK(!buffers.empty()) << "buffer list for " << var << " is empty";
         for (const auto &buffer : buffers) {
           if (!IsReducerV2Buffer(buffer)) {
@@ -1178,6 +1176,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer {
   }
 
   void VisitExpr_(const BufferLoadNode *op) final {
+    accessed_buffers_.insert(op->buffer);
     // Collect buffer from BufferLoad
     if (op->buffer.defined() && op->buffer->data.defined()) {
       if (buffer_data_to_buffers_.count(op->buffer->data)) {
@@ -1208,6 +1207,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer {
   }
 
   void VisitStmt_(const BufferStoreNode *op) final {
+    accessed_buffers_.insert(op->buffer);
     // Collect buffer from BufferStore
     if (op->buffer.defined() && op->buffer->data.defined()) {
       if (buffer_data_to_buffers_.count(op->buffer->data)) {
@@ -1336,6 +1336,7 @@ class BufferUseDefCollector : public IRVisitorWithAnalyzer {
   }
 
   Map<Var, Array<Buffer>> buffer_data_to_buffers_;
+  std::unordered_set<Buffer, ObjectPtrHash, ObjectPtrEqual> accessed_buffers_;
   // Map from Bind variable to its bound expression
   Map<Var, PrimExpr> bind_var_to_expr_;
   std::vector<ObjectRef> infer_list_stmt_;