diff --git a/projects/hipblaslt/tensilelite/Tensile/KernelWriterConversion.py b/projects/hipblaslt/tensilelite/Tensile/KernelWriterConversion.py index fa7fa8145f03..ec53e630d0c3 100644 --- a/projects/hipblaslt/tensilelite/Tensile/KernelWriterConversion.py +++ b/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 diff --git a/projects/hipblaslt/tensilelite/Tensile/Source/memory_gfx.h b/projects/hipblaslt/tensilelite/Tensile/Source/memory_gfx.h index f0a8d12689f3..beeb90043b02 100644 --- a/projects/hipblaslt/tensilelite/Tensile/Source/memory_gfx.h +++ b/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(voffset); + uint64_t base = reinterpret_cast(const_cast(base_ptr)); + base += voffset & 0xFFFFFFFF00000000ull; + return reinterpret_cast(base); +} + //////////////////////////////////////////////////////////////////////////////////////////////////// /// @@ -534,6 +543,21 @@ struct buffer_store 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(&D); + llvm_amdgcn_raw_buffer_store_i8( + data, buffer_rsc, voffset_lo, __builtin_amdgcn_readfirstlane(soffset), cache_op); + } }; template @@ -551,6 +575,21 @@ struct buffer_store 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(&D); + llvm_amdgcn_raw_buffer_store_f16( + data, buffer_rsc, voffset_lo, __builtin_amdgcn_readfirstlane(soffset), cache_op); + } }; template @@ -568,6 +607,21 @@ struct buffer_store 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(&D); + llvm_amdgcn_raw_buffer_store_f32( + data, buffer_rsc, voffset_lo, __builtin_amdgcn_readfirstlane(soffset), cache_op); + } }; template @@ -585,6 +639,21 @@ struct buffer_store 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(&D); + llvm_amdgcn_raw_buffer_store_f32x2( + data, buffer_rsc, voffset_lo, __builtin_amdgcn_readfirstlane(soffset), cache_op); + } }; template @@ -602,6 +671,21 @@ struct buffer_store 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(&D); + llvm_amdgcn_raw_buffer_store_f32x4( + data, buffer_rsc, voffset_lo, __builtin_amdgcn_readfirstlane(soffset), cache_op); + } }; //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/projects/hipblaslt/tensilelite/Tensile/Tests/unit/test_KernelWriterConversion_64bit_offset.py b/projects/hipblaslt/tensilelite/Tensile/Tests/unit/test_KernelWriterConversion_64bit_offset.py new file mode 100644 index 000000000000..682a8838e843 --- /dev/null +++ b/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