chore: import upstream snapshot with attribution

samples/python/quickly_deployable_plugins/oait_kernels.py

@@ -0,0 +1,78 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 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 triton
+import triton.language as tl
+
+@triton.jit
+def add_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    x = tl.load(x_ptr + offsets, mask=mask)
+    tl.store(y_ptr + offsets, x + 1, mask=mask)
+
+
+@triton.jit
+def circ_pad_kernel(
+    # input tensor
+    X,
+    # extra scalar args in between input and output tensors
+    # for kernel signature to be compatible with AOT plugin impl
+    all_pads_0,
+    all_pads_2,
+    all_pads_4,
+    all_pads_6,
+    orig_dims_0,
+    orig_dims_1,
+    orig_dims_2,
+    orig_dims_3,
+    Y_shape_1,
+    Y_shape_2,
+    Y_shape_3,
+    X_len,
+    Y_len,
+    # output tensor
+    Y,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    i = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+
+    mask_y = i < Y_len
+
+    i3 = i % Y_shape_3
+    i2 = (i // Y_shape_3) % Y_shape_2
+    i1 = (i // Y_shape_3 // Y_shape_2) % Y_shape_1
+    i0 = i // Y_shape_3 // Y_shape_2 // Y_shape_1
+
+    j0 = (i0 - all_pads_0 + orig_dims_0) % orig_dims_0
+    j1 = (i1 - all_pads_2 + orig_dims_1) % orig_dims_1
+    j2 = (i2 - all_pads_4 + orig_dims_2) % orig_dims_2
+    j3 = (i3 - all_pads_6 + orig_dims_3) % orig_dims_3
+
+    load_idx = (
+        orig_dims_3 * orig_dims_2 * orig_dims_1 * j0
+        + orig_dims_3 * orig_dims_2 * j1
+        + orig_dims_3 * j2
+        + j3
+    )
+    mask_x = load_idx < X_len
+
+    x = tl.load(X + load_idx, mask=mask_x)
+
+    tl.store(Y + i, x, mask=mask_y)