[Bug] [TIRx][CUDA] MergeSharedMemoryAllocations hoists a dynamic allocation above its size binding

オープン
  • 何か問題がありますか？
#20,870 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 1 日以内に返信

まだ誰も着手していません。

評価

難易度
4/5
見積もり時間
3〜5日
初心者へのやさしさ
22/100
issue の種類
バグ
明瞭さ
明確に書かれている
活発さ
活発
技術スタック
cpp, python
領域
compilers

調査の方向性

Start with the reproducer script, which shows the buffer extent size being hoisted above its T.let definition by MergeSharedMemoryAllocations. The use-before-definition check is the assertion at src/tirx/analysis/var_use_def_analysis.cc:127, reached through SplitHostDevice. Done means the merged module keeps size defined before its use and the script runs to passed.

索引モデルが issue の本文から書いたものです。

説明

needs-triage type: bug
Expected behavior

s_tir.transform.MergeSharedMemoryAllocations() should preserve the scope and definition order of variables used in the merged buffer extent. The resulting module should be accepted by tirx.transform.SplitHostDevice().

The example below has two dynamic shared-memory buffers. The second buffer uses a local size binding. The input passes both verify_ssa and verify_well_formed. The intended runtime domain is 1 <= n <= 128.

Actual behavior

The merge pass moves the combined buffer allocation before the definition of size. The transformed IR contains this ordering:

size = T.int64()  # Variable declaration, not a value binding.
buf_dyn_shmem = T.alloc_buffer(
    (T.int64(512) + T.max(T.int64(0), size * T.int64(4)),),
    "uint8",
    scope="shared.dyn",
)
# Buffer declarations omitted.
size = T.min(n, T.int64(128))

SplitHostDevice then fails:

src/tirx/analysis/var_use_def_analysis.cc:127
TVM_FFI_ICHECK(!use_count_.count(v))

tvm.error.InternalError: Check failed: (!use_count_.count(v)) is false:
variable size has been used before definition!

The script below reproduced the same assertion in three fresh processes. Replacing only the second buffer's extent (size,) with the equivalent inline expression (T.min(n, T.int64(128)),) allowed the merge and host/device split to pass in three fresh processes.

This small example does not invoke CSE or Relax compilation. In an FP16 symbolic matmul compiled for CUDA, the same ordering problem was observed between the IR immediately before and after MergeSharedMemoryAllocations, followed by a use-before-definition assertion for cse_v46 in SplitHostDevice.

Environment
  • Apache TVM 0.26.dev0+source.8f328e8 (source snapshot; no .git directory)
  • Ubuntu 22.04.5, Linux 6.8.0-52-generic, x86-64
  • Python 3.11.16, LLVM 15.0.7
  • NVIDIA RTX A6000, driver 550.120
  • CUDA toolkit 12.4, target sm_86, LLVM host target generic

The reproducer only transforms IR; it does not allocate GPU tensors or execute a kernel.

Steps to reproduce

Save the following script as repro_merge_shared.py and run:

python repro_merge_shared.py

For a passing control, change this line:

C = T.alloc_buffer((size,), "float32", scope="shared.dyn")

to:

C = T.alloc_buffer((T.min(n, T.int64(128)),), "float32", scope="shared.dyn")

import tvm
from tvm import s_tir
from tvm.script import ir as I, tirx as T


@I.ir_module
class Module:
    @T.prim_func(s_tir=True)
    def main(A: T.Buffer((128,), "float32"), n: T.int64):
        T.func_attr({
            "target": T.target({
                "kind": "cuda", "arch": "sm_86",
                "host": {"kind": "llvm", "mcpu": "generic"},
            }),
            "tirx.noalias": True,
        })
        bx = T.launch_thread("blockIdx.x", T.int64(1))
        B = T.alloc_buffer((128,), "float32", scope="shared.dyn")
        size: T.let[T.int64] = T.min(n, T.int64(128))
        C = T.alloc_buffer((size,), "float32", scope="shared.dyn")
        tx = T.launch_thread("threadIdx.x", T.int64(128))
        B[tx] = A[tx]
        if tx < size:
            C[tx] = B[tx]
            A[tx] = C[tx] + T.float32(1)


assert tvm.tirx.analysis.verify_ssa(Module["main"])
assert tvm.tirx.analysis.verify_well_formed(Module["main"])
print("input validation passed", flush=True)
merged = s_tir.transform.MergeSharedMemoryAllocations()(Module)
print(merged.script(), flush=True)
tvm.tirx.transform.SplitHostDevice()(merged)
print("passed", flush=True)
Triage
  • needs-triage
  • type: bug
  • backend:cuda
主要言語
Python
スター
13.9k
フォーク
4k
平均マージ
7時間 59分
マージ済み PR（30日）
159

環境構築

このプロジェクトには開発コンテナ、Dockerfile、コントリビューションガイドがありません。まず README を読み、一般的な手順ははじめてのコントリビューションガイドを参照してください。

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

apache/tvm のほかの issue

apache/tvm の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。