[Ascend] Support exact placement of versioned on-chip buffers

src/ascend/target_utils.h

@@ -6,12 +6,21 @@
 #ifndef TVM_TL_ASCEND_TARGET_UTILS_H_
 #define TVM_TL_ASCEND_TARGET_UTILS_H_
 
+#include <cstdint>
+
 #include <tvm/runtime/data_type.h>
 #include <tvm/target/target.h>
 
 namespace tvm {
 namespace tl {
 
+constexpr int64_t kAscendUBMemoryBytes = 248 * 1024;
+constexpr int64_t kAscendSimtContextBytes = 32 * 1024;
+constexpr int64_t kAscendL1MemoryBytes = 512 * 1024;
+constexpr int64_t kAscendL0ABMemoryBytes = 64 * 1024;
+constexpr int64_t kAscendL0CMemoryBytes = 256 * 1024;
+constexpr int64_t kAscendSSMemoryBytes = 3 * 1024;
+
 bool TargetIsAscend(Target target);
 bool IsAscendVectorizableFP8(DataType dtype);
 

src/ascend/transform/auto_schedule.cc

@@ -63,6 +63,7 @@
 #include "./auto_schedule/task_annotations.h"
 #include "ascend/op/copy.h"
 #include "ascend/op/utils.h"
+#include "ascend/transform/buffer_offsets.h"
 #include "ascend/transform/buffer_version.h"
 #include "backend/common/target_utils.h"
 #include "op/utils.h"
@@ -104,37 +105,34 @@ inline bool HasSimtVF(const tirx::Stmt &stmt) {
 }
 
 inline int64_t GetSharedMemoryLimit(const tirx::Stmt &kernel_body) {
-  constexpr int64_t kPhysicalUnifiedBufferSize = 248 * 1024;
-  constexpr int64_t kSimtVFThreadContextSize = 32 * 1024;
-  return HasSimtVF(kernel_body)
-             ? kPhysicalUnifiedBufferSize - kSimtVFThreadContextSize
-             : kPhysicalUnifiedBufferSize;
+  return HasSimtVF(kernel_body) ? kAscendUBMemoryBytes - kAscendSimtContextBytes
+                                : kAscendUBMemoryBytes;
 }
 
 inline int64_t GetL1MemoryLimit(Target target) {
   if (TargetIsAscend(target)) {
-    return 512 * 1024; // 512 KB L1 cache on Ascend
+    return kAscendL1MemoryBytes;
   }
   return 0;
 }
 
 inline int64_t GetL0CMemoryLimit(Target target) {
   if (TargetIsAscend(target)) {
-    return 256 * 1024; // 256 KB L0C buffer on Ascend
+    return kAscendL0CMemoryBytes;
   }
   return 0;
 }
 
 inline int64_t GetL0AMemoryLimit(Target target) {
   if (TargetIsAscend(target)) {
-    return 64 * 1024; // 64 KB L0A buffer on Ascend
+    return kAscendL0ABMemoryBytes;
   }
   return 0;
 }
 
 inline int64_t GetL0BMemoryLimit(Target target) {
   if (TargetIsAscend(target)) {
-    return 64 * 1024; // 64 KB L0B buffer on Ascend
+    return kAscendL0ABMemoryBytes;
   }
   return 0;
 }
@@ -1404,7 +1402,7 @@ static void ScheduleSingleKernel(const Stmt &kernel_body, Target target,
   unit_builder.SetMemoryLimit("shared.l0a", GetL0AMemoryLimit(target));
   unit_builder.SetMemoryLimit("shared.l0b", GetL0BMemoryLimit(target));
   // One SS pool is shared by AIC and both AIVs, unlike per-AIV UB storage.
-  unit_builder.SetMemoryLimit("shared.ss", 3 * 1024);
+  unit_builder.SetMemoryLimit("shared.ss", kAscendSSMemoryBytes);
   constexpr int64_t kUnlimitedMemory = 1LL << 40;
   for (const String &scope : metadata.unlimit_memory_scopes) {
     unit_builder.SetMemoryLimit(scope, kUnlimitedMemory);
@@ -1450,7 +1448,30 @@ tvm::transform::Pass AutoSchedule() {
               context.root, context.outer_ctx, MakeAutoScheduleExtraInfo);
           BufferVersionMap buffer_version_overrides =
               scheduled_tir.metadata.buffer_versions;
+          BufferOffsets offsets;
+          if (auto value = context.root->annotations.Get(kBufferOffsets)) {
+            offsets = ParseBufferOffsets(value.value());
+            // Validate before AutoSchedule replaces frontend overrides with
+            // inferred counts. A mode-only annotation does not fix storage
+            // size.
+            for (const auto &[storage, offset] : offsets) {
+              auto versions = buffer_version_overrides.Get(storage);
+              ICHECK(versions.has_value() && versions.value() > 0)
+                  << kBufferOffsets << " requires an explicit version count in "
+                  << "T.annotate_buffer_versions for storage " << storage;
+            }
+          }
           ScheduleSingleKernel(context.root->body, target, &scheduled_tir);
+          for (const auto &[storage, offset] : offsets) {
+            int requested = buffer_version_overrides.at(storage);
+            int selected =
+                scheduled_tir.metadata.buffer_versions.Get(storage).value_or(1);
+            ICHECK_EQ(selected, requested)
+                << "Cannot place storage " << storage << ": requested "
+                << requested << " versions but AutoSchedule selected "
+                << selected
+                << "; the buffer must be eligible for the requested versions";
+          }
           WarnOnUnappliedBufferVersionOverrides(
               buffer_version_overrides, scheduled_tir.metadata.buffer_versions);
           return EncodeScheduledTIR(std::move(scheduled_tir));

src/ascend/transform/buffer_offsets.h

@@ -0,0 +1,85 @@
+/*!
+ * \file buffer_offsets.h
+ * \brief Exact byte placement of Ascend on-chip storages.
+ */
+
+#ifndef TVM_TL_ASCEND_TRANSFORM_BUFFER_OFFSETS_H_
+#define TVM_TL_ASCEND_TRANSFORM_BUFFER_OFFSETS_H_
+
+#include <string>
+
+#include <tvm/ffi/container/map.h>
+#include <tvm/tirx/expr.h>
+
+#include "ascend/target_utils.h"
+#include "support/check.h"
+#include "tir/transforms/ir_utils.h"
+
+namespace tvm {
+namespace tl {
+
+constexpr const char *kBufferOffsets = "tl.ascend_buffer_offsets";
+using BufferOffsets = ffi::Map<tirx::Var, Integer>;
+
+struct BufferPlacementInfo {
+  int alignment;
+  int64_t capacity;
+};
+
+inline BufferPlacementInfo GetBufferPlacementInfo(const std::string &scope) {
+  if (scope == "shared.dyn")
+    return {32, kAscendUBMemoryBytes};
+  if (scope == "shared.l1")
+    return {512, kAscendL1MemoryBytes};
+  if (scope == "shared.l0a" || scope == "shared.l0b")
+    return {512, kAscendL0ABMemoryBytes};
+  if (scope == "shared.l0c")
+    return {1024, kAscendL0CMemoryBytes};
+  if (scope == "shared.ss")
+    return {32, kAscendSSMemoryBytes};
+  LOG(FATAL)
+      << kBufferOffsets << " does not support storage scope " << scope
+      << "; expected UB, L1, L0A, L0B, L0C or SS allocations. "
+      << "Place bound L0 scale-factor handles through their data buffer.";
+  return {};
+}
+
+inline BufferOffsets ParseBufferOffsets(const ffi::Any &value) {
+  auto offsets = value.try_cast<BufferOffsets>();
+  ICHECK(offsets.has_value())
+      << kBufferOffsets << " must map storage Vars to constant byte offsets";
+  ICHECK(!offsets.value().empty()) << kBufferOffsets << " must not be empty";
+  for (const auto &[storage, offset] : offsets.value()) {
+    std::string scope = GetPtrStorageScope(storage);
+    auto info = GetBufferPlacementInfo(scope);
+    ICHECK(offset->dtype.is_int() && offset->value >= 0 &&
+           offset->value % info.alignment == 0)
+        << kBufferOffsets << " requires non-negative, " << info.alignment
+        << "-byte aligned integer offsets for " << scope << "; storage "
+        << storage << " has " << offset;
+    ICHECK_LT(offset->value, info.capacity)
+        << "Buffer placement for " << storage << " exceeds " << scope
+        << " capacity of " << info.capacity << " bytes";
+  }
+  return offsets.value();
+}
+
+inline BufferOffsets
+RemapBufferOffsets(const BufferOffsets &offsets,
+                   const ffi::Map<tirx::Var, tirx::Var> &storage_remap) {
+  BufferOffsets result;
+  for (const auto &[storage, offset] : offsets) {
+    tirx::Var remapped = storage_remap.Get(storage).value_or(storage);
+    auto existing = result.Get(remapped);
+    ICHECK(!existing.has_value() || existing.value()->value == offset->value)
+        << "Conflicting " << kBufferOffsets << " after remapping storage "
+        << remapped;
+    result.Set(remapped, offset);
+  }
+  return result;
+}
+
+} // namespace tl
+} // namespace tvm
+
+#endif // TVM_TL_ASCEND_TRANSFORM_BUFFER_OFFSETS_H_

src/ascend/transform/lower_scheduled_tir.cc

@@ -25,6 +25,7 @@
 #include "./auto_schedule/kernel_rewriter.h"
 #include "./auto_schedule/scheduled_tir.h"
 #include "buffer_alias.h"
+#include "buffer_offsets.h"
 
 #include <tvm/ffi/container/array.h>
 #include <tvm/ffi/container/map.h>
@@ -763,6 +764,11 @@ tvm::transform::Pass LowerScheduledTIR() {
                             body);
             writer->annotations.erase(kAivStorageMap);
           }
+          if (auto offsets = root->annotations.Get(kBufferOffsets)) {
+            body = AttrStmt(ParseBufferOffsets(offsets.value()), kBufferOffsets,
+                            Integer(1), std::move(body));
+            writer->annotations.erase(kBufferOffsets);
+          }
           writer->annotations.erase(kSplitSubcores);
           writer->body = std::move(body);
           return root;

src/ascend/transform/lower_tile_op.cc

@@ -38,6 +38,7 @@
 #include "transform/loop_partition.h"
 
 #include "ascend/transform/buffer_alias.h"
+#include "ascend/transform/buffer_offsets.h"
 #include "ascend/transform/buffer_version.h"
 #include "ascend/transform/vf_regions.h"
 
@@ -333,6 +334,9 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer {
   }
 
   Stmt VisitStmt_(const SBlockNode *op) final {
+    ICHECK(!op->annotations.count(kBufferOffsets))
+        << kBufferOffsets << " requires the Ascend scheduled pipeline; "
+        << "enable TL_ENABLE_AUTO_SCHEDULE";
     Map<String, Any> previous_block_annotations = block_annotations_;
     Map<Var, PrimExpr> previous_safe_value_map = safe_value_map_;
     block_annotations_ = op->annotations;
@@ -1279,6 +1283,12 @@ class LowerTileOpPass : arith::IRMutatorWithAnalyzer {
       }
     }
     Stmt stmt = arith::IRMutatorWithAnalyzer::VisitStmt_(op);
+    if (op->attr_key == tl::kBufferOffsets) {
+      AttrStmt attribute = Downcast<AttrStmt>(std::move(stmt));
+      attribute.CopyOnWrite()->node =
+          RemapBufferOffsets(ParseBufferOffsets(op->node), var_remap_);
+      return attribute;
+    }
     if (op->attr_key == tl::kAivStorageMap) {
       AivStorageMap remapped;
       for (const auto &[storage, sid] : op->node.cast<AivStorageMap>()) {