chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Waiting to run
Docker Image CI / build-ubuntu2004 (push) Waiting to run
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,386 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed 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.
|
||||
#
|
||||
|
||||
import tensorrt as trt
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from typing import Tuple, List, Union
|
||||
|
||||
import tensorrt.plugin as trtp
|
||||
import numpy.typing as npt
|
||||
|
||||
import logging
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logging.getLogger("QuicklyDeployablePlugins").setLevel(logging.INFO)
|
||||
|
||||
########## Elemwise-add plugin definition ##########
|
||||
|
||||
|
||||
@trtp.register("sample::elemwise_add_plugin")
|
||||
def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> trtp.TensorDesc:
|
||||
return inp0.like()
|
||||
|
||||
|
||||
# Helper to simulate defining/omitting an autotune definition for the plugin
|
||||
def register_autotune():
|
||||
# Type annotations can be omitted for autotune and impl definitions, but will be checked for consistency if added
|
||||
@trtp.autotune("sample::elemwise_add_plugin")
|
||||
def add_plugin_autotune(
|
||||
inp0: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]
|
||||
) -> List[trtp.AutoTuneCombination]:
|
||||
return [trtp.AutoTuneCombination("FP32|FP16, FP32|FP16")]
|
||||
|
||||
|
||||
@trtp.impl("sample::elemwise_add_plugin")
|
||||
def add_plugin_impl(
|
||||
inp0: trtp.Tensor, block_size: int, outputs: Tuple[trtp.Tensor], stream: int
|
||||
) -> None:
|
||||
|
||||
log = logging.getLogger("QuicklyDeployablePlugins")
|
||||
log.debug(
|
||||
f"Executing for inp0: dtype={inp0.dtype},format={inp0.format} and output[0]: dtype={outputs[0].dtype},format={outputs[0].format}"
|
||||
)
|
||||
|
||||
n = inp0.numel()
|
||||
|
||||
with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
|
||||
inp0_t = torch.as_tensor(inp0, device="cuda")
|
||||
out_t = torch.as_tensor(outputs[0], device="cuda")
|
||||
|
||||
import triton
|
||||
from oait_kernels import add_kernel
|
||||
|
||||
add_kernel[(triton.cdiv(n, block_size),)](inp0_t, out_t, n, BLOCK_SIZE=block_size)
|
||||
|
||||
|
||||
########## In-place elemwise-add plugin definition ##########
|
||||
|
||||
|
||||
@trtp.register("sample::elemwise_add_plugin_")
|
||||
def add_plugin_desc_(inp0: trtp.TensorDesc, delta: int) -> trtp.TensorDesc:
|
||||
return inp0.aliased()
|
||||
|
||||
|
||||
@trtp.autotune("sample::elemwise_add_plugin_")
|
||||
def add_plugin_autotune_(inp0, outputs) -> List[trtp.AutoTuneCombination]:
|
||||
return [
|
||||
trtp.AutoTuneCombination("FP32, FP32", "LINEAR*HWC"),
|
||||
trtp.AutoTuneCombination("FP32|FP16, FP32|FP16", "LINEAR"),
|
||||
]
|
||||
|
||||
|
||||
@trtp.impl("sample::elemwise_add_plugin_")
|
||||
def add_plugin_impl_(inp0, delta: int, outputs, stream) -> None:
|
||||
|
||||
log = logging.getLogger("QuicklyDeployablePlugins")
|
||||
log.debug(
|
||||
f"Executing for inp0: dtype={inp0.dtype},format={inp0.format} and output[0]: dtype={outputs[0].dtype},format={outputs[0].format}"
|
||||
)
|
||||
|
||||
with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
|
||||
inp0_t = torch.as_tensor(inp0, device="cuda")
|
||||
inp0_t.add_(delta)
|
||||
|
||||
|
||||
########## Non-zero plugin (DDS) ##########
|
||||
|
||||
|
||||
@trtp.register("sample::non_zero_plugin")
|
||||
def non_zero_plugin_reg(
|
||||
inp0: trtp.TensorDesc,
|
||||
) -> Tuple[trtp.TensorDesc, trtp.TensorDesc]:
|
||||
upper_bound = inp0.shape_expr[0] * inp0.shape_expr[1]
|
||||
st = trtp.size_tensor(upper_bound // 2, upper_bound)
|
||||
st.dtype = trt.int64
|
||||
return trtp.from_shape_expr((st.expr(), 2), dtype=trt.int32), st
|
||||
|
||||
|
||||
@trtp.autotune("sample::non_zero_plugin")
|
||||
def non_zero_plugin_autotune(inp0, outputs) -> List[trtp.AutoTuneCombination]:
|
||||
return [trtp.AutoTuneCombination("FP32|FP16, INT32, INT64")]
|
||||
|
||||
|
||||
@trtp.impl("sample::non_zero_plugin")
|
||||
def non_zero_plugin_impl(inp0, outputs, stream) -> None:
|
||||
|
||||
log = logging.getLogger("QuicklyDeployablePlugins")
|
||||
log.debug(
|
||||
f"Executing for inp0: dtype={inp0.dtype},format={inp0.format} and output[0]: dtype={outputs[0].dtype},format={outputs[0].format}"
|
||||
)
|
||||
|
||||
with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
|
||||
inp0_t = torch.as_tensor(inp0, device="cuda")
|
||||
out_1 = torch.as_tensor(outputs[1], device="cuda").reshape((-1,))
|
||||
|
||||
out = torch.nonzero(inp0_t)
|
||||
|
||||
out0 = torch.as_tensor(outputs[0].aliased(out.shape), device="cuda")
|
||||
out0.copy_(out)
|
||||
out_1.copy_(torch.Tensor([out.shape[0]]))
|
||||
|
||||
|
||||
########## Circular padding plugin ########
|
||||
|
||||
|
||||
@trtp.register("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_desc(
|
||||
inp0: trtp.TensorDesc, pads: npt.NDArray[np.int32]
|
||||
) -> trtp.TensorDesc:
|
||||
ndim = inp0.ndim
|
||||
out_desc = inp0.like()
|
||||
|
||||
for i in range(np.size(pads) // 2):
|
||||
out_desc.shape_expr[ndim - i - 1] += int(pads[i * 2] + pads[i * 2 + 1])
|
||||
|
||||
return out_desc
|
||||
|
||||
|
||||
# Helper to define a multi-tactic implementation of the plugin
|
||||
def enable_multi_tactic_circ_pad():
|
||||
|
||||
from enum import IntEnum
|
||||
|
||||
class Tactic(IntEnum):
|
||||
TORCH = 1
|
||||
TRITON = 2
|
||||
|
||||
@trtp.autotune("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_autotune(
|
||||
inp0: trtp.TensorDesc,
|
||||
outputs: Tuple[trtp.TensorDesc],
|
||||
) -> List[trtp.AutoTuneCombination]:
|
||||
c = trtp.AutoTuneCombination()
|
||||
c.pos([0, 1], "FP32|FP16")
|
||||
c.tactics([int(Tactic.TORCH), int(Tactic.TRITON)])
|
||||
return [c]
|
||||
|
||||
@trtp.impl("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_impl(
|
||||
inp0: trtp.Tensor,
|
||||
pads: npt.NDArray[np.int32],
|
||||
outputs: Tuple[trtp.Tensor],
|
||||
stream: int,
|
||||
tactic: int,
|
||||
) -> None:
|
||||
|
||||
log = logging.getLogger("QuicklyDeployablePlugins")
|
||||
log.debug(
|
||||
f"Executing for inp0: dtype={inp0.dtype},format={inp0.format} and output[0]: dtype={outputs[0].dtype},format={outputs[0].format}"
|
||||
)
|
||||
|
||||
with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
|
||||
inp_t = torch.as_tensor(inp0, device="cuda")
|
||||
out_t = torch.as_tensor(outputs[0], device="cuda")
|
||||
|
||||
if tactic == Tactic.TORCH:
|
||||
out = torch.nn.functional.pad(inp_t, pads.tolist(), mode="circular")
|
||||
out_t.copy_(out)
|
||||
elif tactic == Tactic.TRITON:
|
||||
N = inp0.ndim
|
||||
all_pads = np.zeros((N * 2,), dtype=np.int32)
|
||||
out_dims = trtp.Shape(tuple(inp0.shape))
|
||||
|
||||
for i in range(np.size(pads) // 2):
|
||||
out_dims[N - i - 1] += pads[i * 2] + pads[i * 2 + 1]
|
||||
all_pads[N * 2 - 2 * i - 2] = pads[i * 2]
|
||||
all_pads[N * 2 - 2 * i - 1] = pads[i * 2 + 1]
|
||||
|
||||
all_pads = all_pads.tolist()
|
||||
|
||||
block_size = 256
|
||||
num_blocks = tuple(
|
||||
[int((np.prod(out_dims) + block_size - 1) // block_size)]
|
||||
)
|
||||
|
||||
from oait_kernels import circ_pad
|
||||
|
||||
circ_pad[num_blocks](
|
||||
inp_t,
|
||||
all_pads[0],
|
||||
all_pads[2],
|
||||
all_pads[4],
|
||||
all_pads[6],
|
||||
inp0.shape[0],
|
||||
inp0.shape[1],
|
||||
inp0.shape[2],
|
||||
inp0.shape[3],
|
||||
int(out_dims[1]),
|
||||
int(out_dims[2]),
|
||||
int(out_dims[3]),
|
||||
inp0.numel(),
|
||||
out_dims.numel(),
|
||||
out_t,
|
||||
BLOCK_SIZE=block_size,
|
||||
)
|
||||
|
||||
|
||||
# Shared AOT compilation body for the circ_pad plugin: build SymInt args,
|
||||
# compile the Triton kernel with the given BLOCK_SIZE, and pack launch params.
|
||||
def _compile_circ_pad_aot(
|
||||
inp0: trtp.TensorDesc,
|
||||
pads: npt.NDArray[np.int32],
|
||||
outputs: Tuple[trtp.TensorDesc],
|
||||
block_size: int,
|
||||
) -> Tuple[Union[str, bytes], Union[str, bytes], trtp.KernelLaunchParams, trtp.SymExprs]:
|
||||
|
||||
N = inp0.ndim
|
||||
all_pads = np.zeros((N * 2,), dtype=np.int32)
|
||||
inp_dims = inp0.shape_expr
|
||||
out_dims = outputs[0].shape_expr
|
||||
|
||||
for i in range(np.size(pads) // 2):
|
||||
all_pads[N * 2 - 2 * i - 2] = pads[i * 2]
|
||||
all_pads[N * 2 - 2 * i - 1] = pads[i * 2 + 1]
|
||||
|
||||
all_pads = all_pads.tolist()
|
||||
|
||||
# Representing all int32-scalar-kernel-inputs as symbolic expressions.
|
||||
# These inputs are either constants or derivatives of input/output shapes (that may be dynamic).
|
||||
# The symbolic expressions are resolved after the full shape context becomes available at runtime.
|
||||
extra_args = trtp.SymIntExprs.from_tuple(
|
||||
[
|
||||
trtp.SymInt32(e)
|
||||
for e in [
|
||||
all_pads[0],
|
||||
all_pads[2],
|
||||
all_pads[4],
|
||||
all_pads[6],
|
||||
inp_dims[0],
|
||||
inp_dims[1],
|
||||
inp_dims[2],
|
||||
inp_dims[3],
|
||||
out_dims[1],
|
||||
out_dims[2],
|
||||
out_dims[3],
|
||||
inp_dims.numel(),
|
||||
out_dims.numel(),
|
||||
]
|
||||
]
|
||||
)
|
||||
|
||||
type_str = "fp32" if inp0.dtype == trt.float32 else "fp16"
|
||||
|
||||
from oait_kernels import circ_pad_kernel
|
||||
import triton
|
||||
|
||||
src = triton.compiler.ASTSource(
|
||||
fn=circ_pad_kernel,
|
||||
signature={
|
||||
"X": f"*{type_str}",
|
||||
"all_pads_0": "i32",
|
||||
"all_pads_2": "i32",
|
||||
"all_pads_4": "i32",
|
||||
"all_pads_6": "i32",
|
||||
"orig_dims_0": "i32",
|
||||
"orig_dims_1": "i32",
|
||||
"orig_dims_2": "i32",
|
||||
"orig_dims_3": "i32",
|
||||
"Y_shape_1": "i32",
|
||||
"Y_shape_2": "i32",
|
||||
"Y_shape_3": "i32",
|
||||
"X_len": "i32",
|
||||
"Y_len": "i32",
|
||||
"Y": f"*{type_str}",
|
||||
},
|
||||
constexprs={"BLOCK_SIZE": block_size},
|
||||
)
|
||||
|
||||
compiled_kernel = triton.compile(src)
|
||||
launch_params = trtp.KernelLaunchParams()
|
||||
launch_params.grid_x = trtp.cdiv(out_dims.numel(), block_size)
|
||||
launch_params.block_x = compiled_kernel.metadata.num_warps * 32
|
||||
launch_params.shared_mem = compiled_kernel.metadata.shared
|
||||
|
||||
return (
|
||||
compiled_kernel.metadata.name.encode(),
|
||||
compiled_kernel.asm["ptx"].encode(),
|
||||
launch_params,
|
||||
extra_args,
|
||||
)
|
||||
|
||||
|
||||
# Helper to define a single tactic implementation of the plugin
|
||||
def enable_single_tactic_circ_pad():
|
||||
@trtp.autotune("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_autotune(
|
||||
inp0: trtp.TensorDesc,
|
||||
outputs: Tuple[trtp.TensorDesc],
|
||||
) -> List[trtp.AutoTuneCombination]:
|
||||
|
||||
return [trtp.AutoTuneCombination("FP32|FP16, FP32|FP16")]
|
||||
|
||||
@trtp.impl("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_impl(
|
||||
inp0: trtp.Tensor,
|
||||
pads: npt.NDArray[np.int32],
|
||||
outputs: Tuple[trtp.Tensor],
|
||||
stream: int,
|
||||
) -> None:
|
||||
with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
|
||||
inp_t = torch.as_tensor(inp0, device="cuda")
|
||||
out_t = torch.as_tensor(outputs[0], device="cuda")
|
||||
|
||||
out = torch.nn.functional.pad(inp_t, pads.tolist(), mode="circular")
|
||||
out_t.copy_(out)
|
||||
|
||||
@trtp.aot_impl("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_aot_impl(
|
||||
inp0: trtp.TensorDesc, pads: npt.NDArray[np.int32], outputs: Tuple[trtp.TensorDesc], tactic: int
|
||||
) -> Tuple[Union[str, bytes], Union[str, bytes], trtp.KernelLaunchParams, trtp.SymExprs]:
|
||||
return _compile_circ_pad_aot(inp0, pads, outputs, block_size=256)
|
||||
|
||||
|
||||
# Helper to define a multi-tactic AOT implementation of the plugin.
|
||||
# Each tactic precompiles the same Triton kernel with a different BLOCK_SIZE,
|
||||
# so TRT times two PTX variants at build time and bakes the winner into the engine.
|
||||
def enable_multi_tactic_aot_circ_pad():
|
||||
|
||||
from enum import IntEnum
|
||||
|
||||
class Tactic(IntEnum):
|
||||
BLOCK_256 = 1
|
||||
BLOCK_1024 = 2
|
||||
|
||||
block_size_by_tactic = {
|
||||
int(Tactic.BLOCK_256): 256,
|
||||
int(Tactic.BLOCK_1024): 1024,
|
||||
}
|
||||
|
||||
@trtp.autotune("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_autotune(
|
||||
inp0: trtp.TensorDesc,
|
||||
outputs: Tuple[trtp.TensorDesc],
|
||||
) -> List[trtp.AutoTuneCombination]:
|
||||
c = trtp.AutoTuneCombination()
|
||||
c.pos([0, 1], "FP32|FP16")
|
||||
c.tactics([int(Tactic.BLOCK_256), int(Tactic.BLOCK_1024)])
|
||||
return [c]
|
||||
|
||||
@trtp.aot_impl("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_aot_impl(
|
||||
inp0: trtp.TensorDesc,
|
||||
pads: npt.NDArray[np.int32],
|
||||
outputs: Tuple[trtp.TensorDesc],
|
||||
tactic: int,
|
||||
) -> Tuple[Union[str, bytes], Union[str, bytes], trtp.KernelLaunchParams, trtp.SymExprs]:
|
||||
block_size = block_size_by_tactic[tactic]
|
||||
logging.getLogger("QuicklyDeployablePlugins").debug(
|
||||
f"aot_impl invoked: tactic={tactic} -> BLOCK_SIZE={block_size}"
|
||||
)
|
||||
return _compile_circ_pad_aot(inp0, pads, outputs, block_size=block_size)
|
||||
Reference in New Issue
Block a user