From 7a25e51b077a7599ce304dbebad9b719b594a8ea Mon Sep 17 00:00:00 2001 From: Caio Rocha <164253795+caiomcbr@users.noreply.github.com> Date: Wed, 23 Apr 2025 16:37:14 -0700 Subject: [PATCH] Automatic creation of Scratch Buffer at MSCCLLang (#510) --- python/mscclpp/language/program.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/python/mscclpp/language/program.py b/python/mscclpp/language/program.py index 1baf0de2..681a2fe3 100644 --- a/python/mscclpp/language/program.py +++ b/python/mscclpp/language/program.py @@ -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)