fix(tensilelite): Fix output-store 32-bit address overflow

projects/hipblaslt/tensilelite/Tensile/KernelWriterConversion.py

@@ -818,15 +818,16 @@ def kernelBody(self):
       kStr += "  %s[%d] = (%s)%s[%d];%s" % (resultStr, vIdx, destTypeStr, accumStr, vIdx, self.endLine)
 
     # kStr += "  *(%s *)(arg.D+idxD) = *(%s *)%s;%s" % (storeTypeStr, storeTypeStr, resultStr, self.endLine)
+    kStr += "  %s byteOffsetD = idxD * sizeof(%s);%s" % (self.uint64Str, destTypeStr, self.endLine)
     if not self.state["ProblemType"]["GroupedGemm"]:
       kStr += "  if(batch_mode == 0) {" + self.endLine
-      kStr += "    buffer_store<%s, sizeof(%s), CacheOperation::Kind::Always>(*(%s *)%s, arg.D, idxD * sizeof(%s), 0);%s" % (storeTypeStr, storeTypeStr, storeTypeStr, resultStr, destTypeStr, self.endLine)
+      kStr += "    buffer_store<%s, sizeof(%s), CacheOperation::Kind::Always>(*(%s *)%s, arg.D, byteOffsetD, 0);%s" % (storeTypeStr, storeTypeStr, storeTypeStr, resultStr, self.endLine)
       kStr += "  } else {" + self.endLine
       kStr += "    %s *ptr = *(reinterpret_cast<%s **>(((char *)arg.D) + (8*id2)));" % (destTypeStr, destTypeStr) + self.endLine
-      kStr += "    buffer_store<%s, sizeof(%s), CacheOperation::Kind::Always>(*(%s *)%s, ptr, idxD * sizeof(%s), 0);%s" % (storeTypeStr, storeTypeStr, storeTypeStr, resultStr, destTypeStr, self.endLine)
+      kStr += "    buffer_store<%s, sizeof(%s), CacheOperation::Kind::Always>(*(%s *)%s, ptr, byteOffsetD, 0);%s" % (storeTypeStr, storeTypeStr, storeTypeStr, resultStr, self.endLine)
       kStr += "  }" + self.endLine
     else:
-      kStr += "    buffer_store<%s, sizeof(%s), CacheOperation::Kind::Always>(*(%s *)%s, arg.D, idxD * sizeof(%s), 0);%s" % (storeTypeStr, storeTypeStr, storeTypeStr, resultStr, destTypeStr, self.endLine)      
+      kStr += "    buffer_store<%s, sizeof(%s), CacheOperation::Kind::Always>(*(%s *)%s, arg.D, byteOffsetD, 0);%s" % (storeTypeStr, storeTypeStr, storeTypeStr, resultStr, self.endLine)
     ########################################
     # end
     kStr += "}%s" % self.endLine

projects/hipblaslt/tensilelite/Tensile/Source/memory_gfx.h

@@ -150,6 +150,15 @@ struct alignas(16) BufferResource
     Desc desc_;
 };
 
+INLINEDEVICE
+void const* splitBufferOffset(void const* base_ptr, uint64_t voffset, uint32_t& voffset_lo)
+{
+    voffset_lo     = static_cast<uint32_t>(voffset);
+    uint64_t base  = reinterpret_cast<uint64_t>(const_cast<void*>(base_ptr));
+    base          += voffset & 0xFFFFFFFF00000000ull;
+    return reinterpret_cast<void const*>(base);
+}
+
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 ///
@@ -534,6 +543,21 @@ struct buffer_store<AccessType, 1, cache_op>
         llvm_amdgcn_raw_buffer_store_i8(
             data, buffer_rsc, voffset, __builtin_amdgcn_readfirstlane(soffset), cache_op);
     }
+
+    INLINEDEVICE
+    buffer_store(const AccessType& D,
+                 void const*       base_ptr,
+                 uint64_t          voffset,
+                 uint32_t          soffset,
+                 uint32_t          num_records = (0xFFFFFFFF - 1))
+    {
+        uint32_t voffset_lo = 0;
+        base_ptr            = splitBufferOffset(base_ptr, voffset, voffset_lo);
+        BufferResource buffer_rsc(base_ptr, num_records);
+        char           data = *reinterpret_cast<char const*>(&D);
+        llvm_amdgcn_raw_buffer_store_i8(
+            data, buffer_rsc, voffset_lo, __builtin_amdgcn_readfirstlane(soffset), cache_op);
+    }
 };
 
 template <typename AccessType, CacheOperation::Kind cache_op>
@@ -551,6 +575,21 @@ struct buffer_store<AccessType, 2, cache_op>
         llvm_amdgcn_raw_buffer_store_f16(
             data, buffer_rsc, voffset, __builtin_amdgcn_readfirstlane(soffset), cache_op);
     }
+
+    INLINEDEVICE
+    buffer_store(const AccessType& D,
+                 void const*       base_ptr,
+                 uint64_t          voffset,
+                 uint32_t          soffset,
+                 uint32_t          num_records = (0xFFFFFFFF - 1))
+    {
+        uint32_t voffset_lo = 0;
+        base_ptr            = splitBufferOffset(base_ptr, voffset, voffset_lo);
+        BufferResource buffer_rsc(base_ptr, num_records);
+        float16_t      data = *reinterpret_cast<float16_t const*>(&D);
+        llvm_amdgcn_raw_buffer_store_f16(
+            data, buffer_rsc, voffset_lo, __builtin_amdgcn_readfirstlane(soffset), cache_op);
+    }
 };
 
 template <typename AccessType, CacheOperation::Kind cache_op>
@@ -568,6 +607,21 @@ struct buffer_store<AccessType, 4, cache_op>
         llvm_amdgcn_raw_buffer_store_f32(
             data, buffer_rsc, voffset, __builtin_amdgcn_readfirstlane(soffset), cache_op);
     }
+
+    INLINEDEVICE
+    buffer_store(const AccessType& D,
+                 void const*       base_ptr,
+                 uint64_t          voffset,
+                 uint32_t          soffset,
+                 uint32_t          num_records = (0xFFFFFFFF - 1))
+    {
+        uint32_t voffset_lo = 0;
+        base_ptr            = splitBufferOffset(base_ptr, voffset, voffset_lo);
+        BufferResource buffer_rsc(base_ptr, num_records);
+        float32_t      data = *reinterpret_cast<float32_t const*>(&D);
+        llvm_amdgcn_raw_buffer_store_f32(
+            data, buffer_rsc, voffset_lo, __builtin_amdgcn_readfirstlane(soffset), cache_op);
+    }
 };
 
 template <typename AccessType, CacheOperation::Kind cache_op>
@@ -585,6 +639,21 @@ struct buffer_store<AccessType, 8, cache_op>
         llvm_amdgcn_raw_buffer_store_f32x2(
             data, buffer_rsc, voffset, __builtin_amdgcn_readfirstlane(soffset), cache_op);
     }
+
+    INLINEDEVICE
+    buffer_store(const AccessType& D,
+                 void const*       base_ptr,
+                 uint64_t          voffset,
+                 uint32_t          soffset,
+                 uint32_t          num_records = (0xFFFFFFFF - 1))
+    {
+        uint32_t voffset_lo = 0;
+        base_ptr            = splitBufferOffset(base_ptr, voffset, voffset_lo);
+        BufferResource buffer_rsc(base_ptr, num_records);
+        float32x2_t    data = *reinterpret_cast<float32x2_t const*>(&D);
+        llvm_amdgcn_raw_buffer_store_f32x2(
+            data, buffer_rsc, voffset_lo, __builtin_amdgcn_readfirstlane(soffset), cache_op);
+    }
 };
 
 template <typename AccessType, CacheOperation::Kind cache_op>
@@ -602,6 +671,21 @@ struct buffer_store<AccessType, 16, cache_op>
         llvm_amdgcn_raw_buffer_store_f32x4(
             data, buffer_rsc, voffset, __builtin_amdgcn_readfirstlane(soffset), cache_op);
     }
+
+    INLINEDEVICE
+    buffer_store(const AccessType& D,
+                 void const*       base_ptr,
+                 uint64_t          voffset,
+                 uint32_t          soffset,
+                 uint32_t          num_records = (0xFFFFFFFF - 1))
+    {
+        uint32_t voffset_lo = 0;
+        base_ptr            = splitBufferOffset(base_ptr, voffset, voffset_lo);
+        BufferResource buffer_rsc(base_ptr, num_records);
+        float32x4_t    data = *reinterpret_cast<float32x4_t const*>(&D);
+        llvm_amdgcn_raw_buffer_store_f32x4(
+            data, buffer_rsc, voffset_lo, __builtin_amdgcn_readfirstlane(soffset), cache_op);
+    }
 };
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////

projects/hipblaslt/tensilelite/Tensile/Tests/unit/test_KernelWriterConversion_64bit_offset.py

@@ -0,0 +1,110 @@
+################################################################################
+#
+# Copyright (C) 2025-2026 Advanced Micro Devices, Inc. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+# SPDX-License-Identifier: MIT
+################################################################################
+"""
+Unit tests for the 64-bit byte offset fix in KernelWriterConversion.
+
+Validates that buffer_store calls in the conversion kernel use a pre-computed
+uint64_t byteOffsetD instead of inline 32-bit index multiplication, preventing
+address overflow for large tensors (>2GB).
+"""
+import re
+import os
+import pytest
+
+
+TENSILE_ROOT = os.path.join(os.path.dirname(__file__), "..", "..")
+
+
+class TestKernelWriterConversion64BitOffset:
+    """Verify KernelWriterConversion uses uint64_t byteOffsetD for buffer_store."""
+
+    @pytest.fixture(autouse=True)
+    def load_source(self):
+        src_path = os.path.join(TENSILE_ROOT, "KernelWriterConversion.py")
+        with open(src_path, "r") as f:
+            self.source = f.read()
+
+    def test_byteOffsetD_declared_as_uint64(self):
+        """byteOffsetD must be declared with self.uint64Str type."""
+        pattern = (
+            r'kStr\s*\+=\s*".*%s byteOffsetD\s*=\s*idxD\s*\*\s*sizeof\(%s\)'
+            r'.*"\s*%\s*\(\s*self\.uint64Str\s*,\s*destTypeStr'
+        )
+        assert re.search(pattern, self.source), \
+            "Expected byteOffsetD declaration to use self.uint64Str"
+
+    def test_buffer_store_uses_byteOffsetD(self):
+        """All buffer_store calls in the store section must use byteOffsetD."""
+        store_lines = [
+            line for line in self.source.splitlines()
+            if "buffer_store" in line and "kStr" in line and "byteOffsetD" in line
+        ]
+        assert len(store_lines) >= 3, \
+            f"Expected at least 3 buffer_store calls using byteOffsetD, found {len(store_lines)}"
+
+    def test_no_inline_idxD_multiply_in_buffer_store(self):
+        """buffer_store must not use inline 'idxD * sizeof(...)' anymore."""
+        store_lines = [
+            line for line in self.source.splitlines()
+            if "buffer_store" in line and "kStr" in line
+        ]
+        for line in store_lines:
+            assert not re.search(r"idxD\s*\*\s*sizeof", line), \
+                f"Found deprecated inline idxD*sizeof in buffer_store: {line.strip()}"
+
+
+class TestMemoryGfxSplitBufferOffset:
+    """Verify memory_gfx.h contains the splitBufferOffset helper and 64-bit overloads."""
+
+    @pytest.fixture(autouse=True)
+    def load_header(self):
+        header_path = os.path.join(TENSILE_ROOT, "Source", "memory_gfx.h")
+        with open(header_path, "r") as f:
+            self.header = f.read()
+
+    def test_splitBufferOffset_exists(self):
+        """splitBufferOffset helper function must exist."""
+        assert "splitBufferOffset" in self.header
+
+    def test_splitBufferOffset_signature(self):
+        """splitBufferOffset must accept base_ptr, uint64_t voffset, and uint32_t& voffset_lo."""
+        pattern = r'splitBufferOffset\s*\(\s*void\s+const\s*\*\s*base_ptr\s*,\s*uint64_t\s+voffset\s*,\s*uint32_t\s*&\s*voffset_lo\s*\)'
+        assert re.search(pattern, self.header), \
+            "splitBufferOffset signature mismatch"
+
+    def test_splitBufferOffset_masks_high_bits(self):
+        """splitBufferOffset must mask off upper 32 bits from voffset into base pointer."""
+        assert "0xFFFFFFFF00000000ull" in self.header
+
+    def test_buffer_store_uint64_overloads_exist(self):
+        """buffer_store specializations must have uint64_t voffset overload."""
+        overload_pattern = r'buffer_store\(.*?uint64_t\s+voffset'
+        matches = re.findall(overload_pattern, self.header, re.DOTALL)
+        assert len(matches) >= 5, \
+            f"Expected at least 5 buffer_store overloads with uint64_t voffset, found {len(matches)}"
+
+    def test_overloads_call_splitBufferOffset(self):
+        """Each uint64_t overload must call splitBufferOffset to split the offset."""
+        assert self.header.count("splitBufferOffset(base_ptr, voffset, voffset_lo)") >= 5