# 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

try:
    from triton.language.extra.hip.libdevice import pow as _pow
except ImportError:
    try:
        from triton.language.math import pow as _pow
    except ImportError:
        from triton.language.libdevice import pow as _pow

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 pow_func_scalar_tensor_wrapper_rank_1(val0, in0: Union[torch.Tensor, StridedBuffer], /, *, out0: Union[torch.Tensor, StridedBuffer]): 
    """Generated wrapper function with Pointwise: scalar, 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):
        pow_func_scalar_tensor_kernel_rank_1[grid](
            val0, 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 pow_func_scalar_tensor(x, exponent):
    return _pow(x.to(tl.float32), exponent)

@triton.jit(do_not_specialize=['val0'])
def pow_func_scalar_tensor_kernel_rank_1(
    val0,
    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 = pow_func_scalar_tensor(val0, 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 = pow_func_scalar_tensor(val0, 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_pow_func_scalar_tensor_wrapper_rank_1():
    # Case 1: Test with a 1D tensor, matching shapes for input and output, scalar exponent
    in_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32, device="cuda")
    out_tensor = torch.zeros_like(in_tensor, device="cuda")
    scalar_value = 2.0  # Exponent for the pow operation
    
    # Call the kernel wrapper function
    pow_func_scalar_tensor_wrapper_rank_1(scalar_value, in_tensor, out0=out_tensor)

    # Case 2: Test with a different exponent (fractional)
    scalar_value = 0.5
    pow_func_scalar_tensor_wrapper_rank_1(scalar_value, in_tensor, out0=out_tensor)

    # Case 3: Test with a different shape (larger tensor)
    in_tensor_large = torch.tensor([2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0], dtype=torch.float32, device="cuda")
    out_tensor_large = torch.zeros_like(in_tensor_large, device="cuda")
    scalar_value = 3.0
    pow_func_scalar_tensor_wrapper_rank_1(scalar_value, in_tensor_large, out0=out_tensor_large)

    # Case 4: Test with negative values in the input tensor
    in_tensor_negative = torch.tensor([-1.0, -2.0, -3.0], dtype=torch.float32, device="cuda")
    out_tensor_negative = torch.zeros_like(in_tensor_negative, device="cuda")
    scalar_value = 2.0  # Exponent should result in positive values for even powers
    pow_func_scalar_tensor_wrapper_rank_1(scalar_value, in_tensor_negative, out0=out_tensor_negative)

    # Case 5: Test with large exponent
    in_tensor_large_exp = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32, device="cuda")
    out_tensor_large_exp = torch.zeros_like(in_tensor_large_exp, device="cuda")
    scalar_value = 10.0  # Large exponent
    pow_func_scalar_tensor_wrapper_rank_1(scalar_value, in_tensor_large_exp, out0=out_tensor_large_exp)

    # Collect results
    results = {
        "test_case_1": out_tensor.clone().cpu().numpy(),
        "test_case_2": out_tensor.clone().cpu().numpy(),
        "test_case_3": out_tensor_large.clone().cpu().numpy(),
        "test_case_4": out_tensor_negative.clone().cpu().numpy(),
        "test_case_5": out_tensor_large_exp.clone().cpu().numpy(),
    }
    return results

# Run the test function
result_gold = test_pow_func_scalar_tensor_wrapper_rank_1()
