Skip to content

[ttl][dfb] A raw_element_read scalar used after the block's pop is rejected in a guarded region #1157

Description

@brnorris03

ttl.raw_element_read returns a scalar loaded from the block when the operation executes, so a use of that scalar after the block's pop does not access the block. ttl-insert-cb-sync nevertheless counts every transitive use of the result as a use of the block's slot, so in a guarded region it rejects a pop that follows the read but precedes a use of the scalar. #1108 stops the slot tracking at ttl.read_index, which returns a scalar the same way (probe R2 is accepted there), but not at ttl.raw_element_read.

To Reproduce

import os

os.environ["TTLANG_COMPILE_ONLY"] = "1"

import torch
import ttl
import ttnn
from ttl import ttl_api

ttl_api._device_target_arch = lambda _runtime_args: "blackhole"


@ttl.operation(grid=(1, 1))
def rawread_after_pop(inp, out):
    d = ttl.make_dataflow_buffer_like(inp, shape=(1, 1), block_count=2)
    e = ttl.make_dataflow_buffer_like(inp, shape=(1, 1), block_count=2)

    @ttl.compute()
    def compute():
        pass

    @ttl.datamovement()
    def reader():
        node_x, _ = ttl.node(dims=2)
        if node_x == 0:
            with d.reserve() as blk:
                ttl.copy(inp[0, 0], blk).wait()

    @ttl.datamovement()
    def writer():
        node_x, _ = ttl.node(dims=2)
        if node_x == 0:
            a = d.wait()
            value = ttl.raw_element_read(a, 0, 0)  # a scalar, loaded here
            a.pop()
            with e.reserve() as blk:
                ttl.raw_element_write(blk, 0, 0, value)  # uses the scalar after the pop
            with e.wait() as blk:
                ttl.copy(blk, out[0, 0]).wait()


tensors = [
    ttnn.from_torch(torch.zeros((32, 32), dtype=torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT)
    for _ in range(2)
]
rawread_after_pop(*tensors)
print("COMPILED")
python repro_rawread_after_pop.py

Observed on main 546f6b20a and with #1108/#1085 (identical):

error: guarded local dataflow buffer pop must follow all uses in its acquiring region
  --> repro_rawread_after_pop.py:34:13
   |
34 |             a.pop()

Expected: the program compiles; the raw read executes before the pop, and the later raw_element_write uses only the loaded scalar. Apply the read_index exclusion (stop slot tracking at an operation that returns a loaded scalar) to ttl.raw_element_read as well.

Found by the review of #1108 (finding S4-4).

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    DFBbugSomething isn't workingcompilerMLIR analysis, verification, transformation, and lowering after the Python frontend

    Type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions