618ZXW

[Triton チュートリアル] triton.jit

Tritonは並列プログラミングのための言語とコンパイラです。カスタムDNN計算カーネルを効率的に記述し、最新のGPUハードウェア上で最大スループットで実行できるようにするためのPythonベースのプログラミング環境を提供するように設計されています。

Triton の中国語ドキュメントの詳細については、→ https://triton.hyper.ai/ をご覧ください。

 triton.jit(fn: T)→ JITFunction[T] triton.jit(*, version=None, repr: Callable | None = None, launch_metadata: Callable | None = None, do_not_specialize: Iterable[int] | None = None, debug: bool | None = None, noinline: bool | None = None)→ Callable[[T], JITFunction[T]]

Triton コンパイラを使用して JIT コンパイルされた関数のデコレータ。

知らせ:

  • JIT コンパイルされた関数が呼び出されると、パラメーターに .data_ptr() メソッドと .dtype 属性がある場合、暗黙的にポインターに変換されます。

知らせ:

  • この関数はGPU上でコンパイルされ、実行されます。以下のものにのみアクセスできます。

    1. Pythonプリミティブ
    2. Triton パッケージの組み込み関数
    3. この関数のパラメータ
    4. その他のJITコンパイル関数

パラメータ:

  • fn (呼び出し可能) - JIT コンパイルされる関数。