Method cache_write

python/tvm/s_tir/schedule/schedule.py:1452–1544

Create a block that reads a buffer region into a write cache. It requires: 1) There is only one block who write the buffer in the scope. 2) The scope block have stage-pipeline property. Parameters ---------- block : SBlockRV | str The producer

(
        self,
        block: SBlockRV | str,
        write_buffer_index: int | str | Buffer,
        storage_scope: str,
        consumer_blocks: list[SBlockRV | str] | None = None,
    )

    @type_checked
    def cache_write(
        self,
        block: SBlockRV | str,
        write_buffer_index: int | str | Buffer,
        storage_scope: str,
        consumer_blocks: list[SBlockRV | str] | None = None,
    ) -> SBlockRV:
        """Create a block that reads a buffer region into a write cache. It requires:


        1) There is only one block who write the buffer in the scope.

        2) The scope block have stage-pipeline property.

        Parameters
        ----------
        block : SBlockRV | str
            The producer block of the target buffer.

        write_buffer_index: int
            The index of the buffer in block's write region, the unique
            name of a write buffer in the block, or a Buffer object
            that is within the blocks write region.

        storage_scope: str
            The target storage scope.

        consumer_blocks: Optional[List[SBlockRV | str]]
            An optional list of consumers that should read directly from the cache.
            If not specified, all consumers will read from the original buffer.

        Returns
        -------
        cached_block : SBlockRV
            The block of the cache stage

        Examples
        --------
        Before cache_write, in TensorIR, the IR is:

        .. code-block:: python

            @T.prim_func(s_tir=True)
            def before_cache_write(a: T.handle, b: T.handle) -> None:
                A = T.match_buffer(a, (128, 128))
                B = T.match_buffer(b, (128, 128))
                for i, j in T.grid(128, 128):
                    with T.sblock("B"):
                        vi, vj = T.axis.remap("SS", [i, j])
                        B[vi, vj] = A[vi, vj] * 2.0

        Create the schedule and cache_write:

        .. code-block:: python

            sch = tvm.s_tir.Schedule(before_cache_write)
            block_b = sch.get_sblock("B")
            sch.cache_write(block_b, 0, "local")