microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-05-13 09:46:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Automatic creation of Scratch Buffer at MSCCLLang (#510)

Browse Source
This commit is contained in:
Caio Rocha
2025-04-23 16:37:14 -07:00
committed by GitHub
parent 474ef0b696
commit 7a25e51b07
1 changed files with 6 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/mscclpp/language/program.py
View File

@@ -431,6 +431,12 @@ class Ref(ChunkRef):
def chunk(rank, buffer, index, size=1) -> Ref:
if buffer is Buffer.scratch:
if buffer not in _curr().buffers[rank]:
_curr().buffers[rank][buffer] = BufferSlice(Buffer.scratch, buffer)
if index >= len(_curr().buffers[rank][buffer]):
_curr().buffers[rank][buffer][index] = ChunkRef(rank, buffer, index, size)
if _curr().buffers[rank][buffer][index] is None:
return None
return _curr().get_ref(rank, buffer, index, size)