wehub-resource-sync
79 lines
2.2 KiB
Python
79 lines
2.2 KiB
Python
#
|
|
# 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)
|