File size: 705 Bytes
7fe6713 61992ce 7fe6713 61992ce 7fe6713 61992ce 7fe6713 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | import nki
import nki.language as nl
import nki.isa as nisa
@nki.jit(platform_target="trn2")
def relu(x):
# Check the first dimension's size to ensure it does not exceed on-chip
# memory tile size, since this simple kernel does not tile inputs.
assert x.shape[0] <= nl.tile_size.pmax
x_tile = sbuf.view(dtype=x.dtype, shape=x.shape)
nisa.dma_copy(dst=x_tile, src=x)
out_tile = sbuf.view(dtype=x.dtype, shape=x.shape)
nisa.tensor_scalar(dst=out_tile, data=x_tile, operand0=0, op0=nl.maximum)
c_output = hbm.view(dtype=x.dtype, shape=x.shape)
nisa.dma_copy(dst=c_output, src=out_tile)
return c_output
from . import layers
__all__ = [
"layers",
"relu",
]
|