wehub-resource-sync
94 lines
3.6 KiB
ReStructuredText
94 lines
3.6 KiB
ReStructuredText
.. Licensed to the Apache Software Foundation (ASF) under one
|
|
or more contributor license agreements. See the NOTICE file
|
|
distributed with this work for additional information
|
|
regarding copyright ownership. The ASF licenses this file
|
|
to you under the Apache License, Version 2.0 (the
|
|
"License"); you may not use this file except in compliance
|
|
with the License. You may obtain a copy of the License at
|
|
|
|
.. http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
.. Unless required by applicable law or agreed to in writing,
|
|
software distributed under the License is distributed on an
|
|
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
KIND, either express or implied. See the License for the
|
|
specific language governing permissions and limitations
|
|
under the License.
|
|
|
|
Your first kernel
|
|
=================
|
|
|
|
A complete write → compile → run → inspect loop. This kernel scales a vector by 2
|
|
with one block of 256 threads.
|
|
|
|
.. code-block:: python
|
|
|
|
import numpy as np
|
|
import tvm
|
|
from tvm.script import tirx as T
|
|
|
|
@T.prim_func
|
|
def scale(A_ptr: T.handle, B_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, (256,), "float32")
|
|
B = T.match_buffer(B_ptr, (256,), "float32")
|
|
|
|
T.device_entry() # everything below runs on the device
|
|
bx = T.cta_id([1]) # 1 block (blockIdx)
|
|
tx = T.thread_id([256]) # 256 threads per block (threadIdx)
|
|
|
|
B[tx] = A[tx] * T.float32(2.0)
|
|
|
|
# compile for CUDA through the TIRx pipeline -> an Executable
|
|
exe = tvm.compile(tvm.IRModule({"main": scale}),
|
|
target=tvm.target.Target("cuda"), tir_pipeline="tirx")
|
|
|
|
dev = tvm.cuda(0)
|
|
a = tvm.runtime.tensor(np.random.rand(256).astype("float32"), device=dev)
|
|
b = tvm.runtime.tensor(np.zeros(256, "float32"), device=dev)
|
|
exe(a, b) # run
|
|
print(exe.mod.imports[0].inspect_source()) # the generated CUDA C
|
|
|
|
What the surrounding calls do:
|
|
|
|
- ``tvm.IRModule({"main": scale})`` wraps the ``PrimFunc`` into an *IRModule* — a
|
|
named collection of functions; ``"main"`` is the entry point the compiler builds
|
|
and the symbol you call.
|
|
- ``tvm.compile(mod, target=..., tir_pipeline="tirx")`` returns an **Executable**:
|
|
the compiled host launcher plus device code, bundled together. You call it like
|
|
a function (``exe(a, b)``); arguments are positional and match the kernel's
|
|
parameter order.
|
|
- ``tvm.cuda(0)`` is a handle to CUDA device 0; ``tvm.runtime.tensor(arr,
|
|
device=dev)`` places data on that device as a TVM tensor (an ``NDArray``).
|
|
- ``exe.mod`` is the underlying runtime ``Module``; ``exe.mod.imports[0]`` is the
|
|
imported device (CUDA) module, and its ``.inspect_source()`` returns the
|
|
generated CUDA C.
|
|
|
|
For this kernel, that last ``print`` produces (boilerplate elided):
|
|
|
|
.. code-block:: c++
|
|
|
|
extern "C" __global__ void __launch_bounds__(256)
|
|
scale_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr) {
|
|
int tx = ((int)threadIdx.x);
|
|
B_ptr[tx] = A_ptr[tx] * 2.0f;
|
|
}
|
|
|
|
Every thread writes one element — a direct map from ``B[tx] = A[tx] * 2.0`` to the
|
|
generated indexing.
|
|
|
|
.. note::
|
|
|
|
The compiled ``Executable`` also accepts CUDA ``torch`` tensors **directly**
|
|
(zero-copy, via DLPack) — no conversion step needed:
|
|
|
|
.. code-block:: python
|
|
|
|
import torch
|
|
a = torch.rand(256, device="cuda")
|
|
b = torch.empty(256, device="cuda")
|
|
exe(a, b)
|
|
torch.testing.assert_close(b, a * 2)
|
|
|
|
The following chapters expand each piece: :doc:`functions`, :doc:`buffers`,
|
|
:doc:`control_flow`, :doc:`threads_sync`, and :doc:`compiling`.
|