add load share offset from load indices

8443d88e · wangziyang · 74e57416 · 8443d88e
Commit 8443d88e authored Apr 16, 2026 by wangziyang
Hide whitespace changes
Inline Side-by-side

Showing with 8 additions and 1 deletion

src/transform/inject_ds_read.cc src/transform/inject_ds_read.cc +8 -1

No files found.
--- a/src/transform/inject_ds_read.cc
+++ b/src/transform/inject_ds_read.cc
@@ -103,7 +103,14 @@ class DSReadInjector : public StmtExprMutator {
    // For A_shared, use the actual shared memory base pointer
    PrimExpr m = make_const(DataType::Int(32), 32);
    PrimExpr n = make_const(DataType::Int(32), 16);
-    PrimExpr offset = make_const(DataType::Int(32), 0);
+
+    PrimExpr offset;
+    // Extract the shared memory offset from the load indices
+    if (!load->indices.empty()) {
+      offset = load->indices[0];
+    } else {
+      offset = make_const(DataType::Int(32), 0);
+    }

    // Use buffer data vars directly
    Array<PrimExpr> new_args = {