# Modifications Copyright(C)[2025] Advanced Micro Devices, Inc. All rights reserved.
# https://github.com/thunlp/TritonBench - Apache License 2.0
import math
from typing import Union
import torch
import triton
from triton import language as tl

def heuristics_for_tile_size(max_tile_size, *sizes):
    ndim = len(sizes)
    tile_sizes = [0 for _ in range(ndim)]
    for i in range(ndim):
        size = sizes[ndim - 1 - i]
        tile_size = min(max_tile_size, triton.next_power_of_2(size))
        tile_sizes[ndim - 1 - i] = tile_size
        max_tile_size = max(1, max_tile_size // tile_size)
    return tuple(tile_sizes)

def heuristics_for_num_warps(tile_size):
    if tile_size < 2048:
        return 4
    elif tile_size < 4096:
        return 8
    else:
        return 16
    

class StridedBuffer:
    """A drop-in replacement of torch.Tensor that can be used in wrapper generated by
    PointwiseDynamicFunction. It allows us to use a different shape, stride, data
    pointer that that of the base tensor.

    It is a kind of reinterpretation of the base tensor. We make this class since we
    cannot get a Tensor view with negative strides via torch APIs, while we need this
    to implement flip op.

    Although generated code can accept torch.Tensor & StridedBuffer, but StridedBuffer
    may not have all the methods as torch.Tensors do. We add some attributes & methods
    with the same name as torch.Tensor, which are used in the generated code. But we
    may not cover all the methods, add one if what you need is missing here.

    And can also be used in triton kernels since it also has dtype & data_ptr().
    """

    def __init__(
        self, base: torch.Tensor, shape=None, strides=None, dtype=None, offset=0
    ):
        self._base = base
        self.dtype = dtype or base.dtype
        if offset == 0:
            self._data_ptr = self._base.data_ptr()
        else:
            offset = self.dtype.itemsize * offset
            self._data_ptr = self._base.data_ptr() + offset
        self.shape = tuple(shape if shape is not None else self._base.shape)
        self._strides = tuple(strides if strides is not None else self._base.stride())
        self.device = self._base.device
        self.ndim = len(self.shape)

    def stride(self):
        return self._strides

    def size(self):
        return self.shape

    def element_size(self):
        return self.dtype.itemsize

    def numel(self):
        return math.prod(self.shape)

    def dim(self):
        return self.ndim

    def unwrap(self):
        return self._base

    def data_ptr(self):
        return self._data_ptr



def relu_forward_wrapper_rank_1(in0: Union[torch.Tensor, StridedBuffer], /, *, out0: Union[torch.Tensor, StridedBuffer]): 
    """Generated wrapper function with Pointwise: StridedBuffer, StridedBuffer(a1!) -> StridedBuffer(a1!)"""
    assert in0.shape == out0.shape, 'operand shapes mismatch'
    # task partitioning
    shape = out0.shape
    num_tasks = out0.numel()
    tile_sizes = heuristics_for_tile_size(512, *shape)
    tile_size = math.prod(tile_sizes)
    num_tiles = math.prod(triton.cdiv(size, tile_size) for size, tile_size in zip(shape, tile_sizes))
    num_ctas = min(65536, num_tiles)
    tiles_per_cta = triton.cdiv(num_tiles, num_ctas)
    num_warps = heuristics_for_num_warps(tile_size)
    one_tile_per_cta = tiles_per_cta==1
    grid = (num_ctas, 1, 1)
    # kernel launch
    in0_strides = in0.stride()
    in0_stride_order = (0,)
    out0_strides = out0.stride()
    out0_stride_order = (0,)
    with torch.cuda._DeviceGuard(in0.device.index):
        relu_forward_kernel_rank_1[grid](
            in0, out0,
            in0_strides[0], # stride for in0
            in0_stride_order[0], # stride order for in0
            out0_strides[0], # stride for out0
            out0_stride_order[0], # stride orderfor out0
            shape[0], # task indexing space
            num_tasks, # num tasks
            tiles_per_cta=tiles_per_cta, # tiles_per_cta
            tile_size0=tile_sizes[0],
            one_tile_per_cta=one_tile_per_cta,
            num_warps=num_warps,
        )
    return out0

@triton.jit
def relu_forward(x):
    return tl.where(x > 0, x, 0)

@triton.jit
def relu_forward_kernel_rank_1(
    in0_ptr: tl.tensor, # of tl.pointer_type
    out0_ptr: tl.tensor, # of tl.pointer_type
    in0_stride0: int, # strides for in0
    in0_stride_order0: tl.constexpr, # stride order for in0
    out0_stride0: int, # strides for out0
    out0_stride_order0: tl.constexpr, # stride order for out0
    s0: int, # task_space
    num_tasks: int,
    tiles_per_cta: int,
    tile_size0: tl.constexpr,
    one_tile_per_cta: tl.constexpr,
):
    pid = tl.program_id(0)
    num_tiles0 = tl.cdiv(s0, tile_size0)
    if one_tile_per_cta: # monolitic kernel style
        tile_id = pid
        # pid multi index recontruction: we use c ordering, right axes changes fastest
        tile_id0 = tile_id

        # tile offsets
        offset0 = tile_id0 * tile_size0
        # loads
        in0_bptr = tl.make_block_ptr(in0_ptr, (s0,), (in0_stride0,), (offset0,), (tile_size0,), order=(in0_stride_order0,))
        in0 = tl.load(in0_bptr, boundary_check=(in0_stride_order0,)).to(in0_ptr.type.element_ty) # workaround the bug on bool, we should use the original pointer's dtype(instead of block pointer's)

        # compute
        out0 = relu_forward(in0)

        # stores, note that store to block pointer does not automatically cast the value to the pointer's dtype
        out0_bptr = tl.make_block_ptr(out0_ptr, (s0,), (out0_stride0,), (offset0,), (tile_size0,), order=(out0_stride_order0,))
        tl.store(out0_bptr, out0.to(out0_bptr.type.element_ty), boundary_check=(out0_stride_order0,))
    else: # grid-stride-loop style kernel
        num_ctas = tl.num_programs(0)
        for j in range(0, tiles_per_cta):
            tile_id = pid + j * num_ctas
            # pid multi index recontruction: we use c ordering, right axes changes fastest
            tile_id0 = tile_id

            # tile offsets
            offset0 = tile_id0 * tile_size0
            # loads
            in0_bptr = tl.make_block_ptr(in0_ptr, (s0,), (in0_stride0,), (offset0,), (tile_size0,), order=(in0_stride_order0,))
            in0 = tl.load(in0_bptr, boundary_check=(in0_stride_order0,)).to(in0_ptr.type.element_ty) # workaround the bug on bool, we should use the original pointer's dtype(instead of block pointer's)

            # compute
            out0 = relu_forward(in0)

            # stores, note that store to block pointer does not automatically cast the value to the pointer's dtype
            out0_bptr = tl.make_block_ptr(out0_ptr, (s0,), (out0_stride0,), (offset0,), (tile_size0,), order=(out0_stride_order0,))
            tl.store(out0_bptr, out0.to(out0_bptr.type.element_ty), boundary_check=(out0_stride_order0,))




##################################################################################################################################################


def test_relu_forward():
    # 测试用的装置 (CUDA)
    device = torch.device('cuda')

    results = {}

    # Test 1: 输入是1维张量，尺寸刚好为tile size的倍数 (简单场景)
    in0 = torch.randn(512, device=device)
    out0 = torch.empty_like(in0)
    relu_forward_wrapper_rank_1(in0, out0=out0)
    results['test_case_1'] = out0

    # Test 2: 输入是1维张量，尺寸小于tile size (小输入)
    in0 = torch.randn(100, device=device)
    out0 = torch.empty_like(in0)
    relu_forward_wrapper_rank_1(in0, out0=out0)
    results['test_case_2'] = out0

    # Test 3: 输入是1维张量，尺寸大于tile size但不是倍数 (复杂大小)
    in0 = torch.randn(1025, device=device)
    out0 = torch.empty_like(in0)
    relu_forward_wrapper_rank_1(in0, out0=out0)
    results['test_case_3'] = out0

    # Test 4: 边界测试，输入维度接近边界大小 (4096)
    in0 = torch.randn(4096, device=device)
    out0 = torch.empty_like(in0)
    relu_forward_wrapper_rank_1(in0, out0=out0)
    results['test_case_4'] = out0

    # Test 5: 测试超大输入张量
    in0 = torch.randn(10000, device=device)
    out0 = torch.empty_like(in0)
    relu_forward_wrapper_rank_1(in0, out0=out0)
    results['test_case_5'] = out0

    # Test 6: 使用 StridedBuffer 的张量操作
    base = torch.randn(512, device=device)
    shape = (512,)
    strides = (1,)
    strided_buffer = StridedBuffer(base, shape=shape, strides=strides, dtype=base.dtype)
    out0 = torch.empty_like(base)
    relu_forward_wrapper_rank_1(strided_buffer, out0=out0)
    results['test_case_6'] = out0

    return results

# 运行测试
result_gold = test_relu_forward()
