import torch
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
from torch._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import math as tl_math
import functools
import numpy as np
import torch.nn.functional as F
import torch.nn as nn
import torch._C
import torch.serialization
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda


@triton.jit
def triton_poi_fused__softmax_0(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr
    ):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp3 = triton_helpers.maximum(tmp1, tmp2)
    tmp5 = triton_helpers.maximum(tmp3, tmp4)
    tmp7 = triton_helpers.maximum(tmp5, tmp6)
    tmp8 = tmp0 - tmp7
    tmp9 = tl_math.exp(tmp8)
    tl.store(out_ptr0 + x2, tmp9, xmask)


@triton.jit
def triton_poi_fused__softmax_1(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr
    ):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp3 = tmp1 + tmp2
    tmp5 = tmp3 + tmp4
    tmp7 = tmp5 + tmp6
    tmp8 = tmp0 / tmp7
    tl.store(out_ptr0 + x2, tmp8, xmask)


@triton.jit
def triton_per_fused__to_copy_add_div_mean_mul_ne_pow_rsub_sum_view_2(
    in_out_ptr0, in_ptr0, in_ptr1, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + 4 * r0, None, eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr1 + 4 * r0, None, eviction_policy='evict_last')
    tmp16 = tl.load(in_ptr1 + (1 + 4 * r0), None, eviction_policy='evict_last')
    tmp29 = tl.load(in_ptr1 + (2 + 4 * r0), None, eviction_policy='evict_last')
    tmp42 = tl.load(in_ptr1 + (3 + 4 * r0), None, eviction_policy='evict_last')
    tmp71 = tl.load(in_ptr0 + (1 + 4 * r0), None, eviction_policy='evict_last')
    tmp112 = tl.load(in_ptr0 + (2 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp153 = tl.load(in_ptr0 + (3 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp2 = tmp1.to(tl.int64)
    tmp3 = tl.full([1, 1], 0, tl.int64)
    tmp4 = triton_helpers.maximum(tmp2, tmp3)
    tmp5 = tl.full([1, 1], 3, tl.int64)
    tmp6 = triton_helpers.minimum(tmp4, tmp5)
    tmp7 = tmp6 == tmp3
    tmp8 = tmp7.to(tl.int64)
    tmp9 = tmp8.to(tl.float32)
    tmp10 = tmp0 * tmp9
    tmp11 = 255.0
    tmp12 = tmp1 != tmp11
    tmp13 = tmp12.to(tl.int64)
    tmp14 = tmp13.to(tl.float32)
    tmp15 = tmp10 * tmp14
    tmp17 = tmp16.to(tl.int64)
    tmp18 = triton_helpers.maximum(tmp17, tmp3)
    tmp19 = triton_helpers.minimum(tmp18, tmp5)
    tmp20 = tmp19 == tmp3
    tmp21 = tmp20.to(tl.int64)
    tmp22 = tmp21.to(tl.float32)
    tmp23 = tmp0 * tmp22
    tmp24 = tmp16 != tmp11
    tmp25 = tmp24.to(tl.int64)
    tmp26 = tmp25.to(tl.float32)
    tmp27 = tmp23 * tmp26
    tmp28 = tmp15 + tmp27
    tmp30 = tmp29.to(tl.int64)
    tmp31 = triton_helpers.maximum(tmp30, tmp3)
    tmp32 = triton_helpers.minimum(tmp31, tmp5)
    tmp33 = tmp32 == tmp3
    tmp34 = tmp33.to(tl.int64)
    tmp35 = tmp34.to(tl.float32)
    tmp36 = tmp0 * tmp35
    tmp37 = tmp29 != tmp11
    tmp38 = tmp37.to(tl.int64)
    tmp39 = tmp38.to(tl.float32)
    tmp40 = tmp36 * tmp39
    tmp41 = tmp28 + tmp40
    tmp43 = tmp42.to(tl.int64)
    tmp44 = triton_helpers.maximum(tmp43, tmp3)
    tmp45 = triton_helpers.minimum(tmp44, tmp5)
    tmp46 = tmp45 == tmp3
    tmp47 = tmp46.to(tl.int64)
    tmp48 = tmp47.to(tl.float32)
    tmp49 = tmp0 * tmp48
    tmp50 = tmp42 != tmp11
    tmp51 = tmp50.to(tl.int64)
    tmp52 = tmp51.to(tl.float32)
    tmp53 = tmp49 * tmp52
    tmp54 = tmp41 + tmp53
    tmp55 = tmp0 * tmp0
    tmp56 = tmp8 * tmp8
    tmp57 = tmp56.to(tl.float32)
    tmp58 = tmp55 + tmp57
    tmp59 = tmp21 * tmp21
    tmp60 = tmp59.to(tl.float32)
    tmp61 = tmp55 + tmp60
    tmp62 = tmp58 + tmp61
    tmp63 = tmp34 * tmp34
    tmp64 = tmp63.to(tl.float32)
    tmp65 = tmp55 + tmp64
    tmp66 = tmp62 + tmp65
    tmp67 = tmp47 * tmp47
    tmp68 = tmp67.to(tl.float32)
    tmp69 = tmp55 + tmp68
    tmp70 = tmp66 + tmp69
    tmp72 = tl.full([1, 1], 1, tl.int64)
    tmp73 = tmp6 == tmp72
    tmp74 = tmp73.to(tl.int64)
    tmp75 = tmp74.to(tl.float32)
    tmp76 = tmp71 * tmp75
    tmp77 = tmp76 * tmp14
    tmp78 = tmp19 == tmp72
    tmp79 = tmp78.to(tl.int64)
    tmp80 = tmp79.to(tl.float32)
    tmp81 = tmp71 * tmp80
    tmp82 = tmp81 * tmp26
    tmp83 = tmp77 + tmp82
    tmp84 = tmp32 == tmp72
    tmp85 = tmp84.to(tl.int64)
    tmp86 = tmp85.to(tl.float32)
    tmp87 = tmp71 * tmp86
    tmp88 = tmp87 * tmp39
    tmp89 = tmp83 + tmp88
    tmp90 = tmp45 == tmp72
    tmp91 = tmp90.to(tl.int64)
    tmp92 = tmp91.to(tl.float32)
    tmp93 = tmp71 * tmp92
    tmp94 = tmp93 * tmp52
    tmp95 = tmp89 + tmp94
    tmp96 = tmp71 * tmp71
    tmp97 = tmp74 * tmp74
    tmp98 = tmp97.to(tl.float32)
    tmp99 = tmp96 + tmp98
    tmp100 = tmp79 * tmp79
    tmp101 = tmp100.to(tl.float32)
    tmp102 = tmp96 + tmp101
    tmp103 = tmp99 + tmp102
    tmp104 = tmp85 * tmp85
    tmp105 = tmp104.to(tl.float32)
    tmp106 = tmp96 + tmp105
    tmp107 = tmp103 + tmp106
    tmp108 = tmp91 * tmp91
    tmp109 = tmp108.to(tl.float32)
    tmp110 = tmp96 + tmp109
    tmp111 = tmp107 + tmp110
    tmp113 = tl.full([1, 1], 2, tl.int64)
    tmp114 = tmp6 == tmp113
    tmp115 = tmp114.to(tl.int64)
    tmp116 = tmp115.to(tl.float32)
    tmp117 = tmp112 * tmp116
    tmp118 = tmp117 * tmp14
    tmp119 = tmp19 == tmp113
    tmp120 = tmp119.to(tl.int64)
    tmp121 = tmp120.to(tl.float32)
    tmp122 = tmp112 * tmp121
    tmp123 = tmp122 * tmp26
    tmp124 = tmp118 + tmp123
    tmp125 = tmp32 == tmp113
    tmp126 = tmp125.to(tl.int64)
    tmp127 = tmp126.to(tl.float32)
    tmp128 = tmp112 * tmp127
    tmp129 = tmp128 * tmp39
    tmp130 = tmp124 + tmp129
    tmp131 = tmp45 == tmp113
    tmp132 = tmp131.to(tl.int64)
    tmp133 = tmp132.to(tl.float32)
    tmp134 = tmp112 * tmp133
    tmp135 = tmp134 * tmp52
    tmp136 = tmp130 + tmp135
    tmp137 = tmp112 * tmp112
    tmp138 = tmp115 * tmp115
    tmp139 = tmp138.to(tl.float32)
    tmp140 = tmp137 + tmp139
    tmp141 = tmp120 * tmp120
    tmp142 = tmp141.to(tl.float32)
    tmp143 = tmp137 + tmp142
    tmp144 = tmp140 + tmp143
    tmp145 = tmp126 * tmp126
    tmp146 = tmp145.to(tl.float32)
    tmp147 = tmp137 + tmp146
    tmp148 = tmp144 + tmp147
    tmp149 = tmp132 * tmp132
    tmp150 = tmp149.to(tl.float32)
    tmp151 = tmp137 + tmp150
    tmp152 = tmp148 + tmp151
    tmp154 = tmp6 == tmp5
    tmp155 = tmp154.to(tl.int64)
    tmp156 = tmp155.to(tl.float32)
    tmp157 = tmp153 * tmp156
    tmp158 = tmp157 * tmp14
    tmp159 = tmp19 == tmp5
    tmp160 = tmp159.to(tl.int64)
    tmp161 = tmp160.to(tl.float32)
    tmp162 = tmp153 * tmp161
    tmp163 = tmp162 * tmp26
    tmp164 = tmp158 + tmp163
    tmp165 = tmp32 == tmp5
    tmp166 = tmp165.to(tl.int64)
    tmp167 = tmp166.to(tl.float32)
    tmp168 = tmp153 * tmp167
    tmp169 = tmp168 * tmp39
    tmp170 = tmp164 + tmp169
    tmp171 = tmp45 == tmp5
    tmp172 = tmp171.to(tl.int64)
    tmp173 = tmp172.to(tl.float32)
    tmp174 = tmp153 * tmp173
    tmp175 = tmp174 * tmp52
    tmp176 = tmp170 + tmp175
    tmp177 = tmp153 * tmp153
    tmp178 = tmp155 * tmp155
    tmp179 = tmp178.to(tl.float32)
    tmp180 = tmp177 + tmp179
    tmp181 = tmp160 * tmp160
    tmp182 = tmp181.to(tl.float32)
    tmp183 = tmp177 + tmp182
    tmp184 = tmp180 + tmp183
    tmp185 = tmp166 * tmp166
    tmp186 = tmp185.to(tl.float32)
    tmp187 = tmp177 + tmp186
    tmp188 = tmp184 + tmp187
    tmp189 = tmp172 * tmp172
    tmp190 = tmp189.to(tl.float32)
    tmp191 = tmp177 + tmp190
    tmp192 = tmp188 + tmp191
    tmp193 = 2.0
    tmp194 = tmp54 * tmp193
    tmp195 = 1.0
    tmp196 = tmp194 + tmp195
    tmp197 = tmp70 + tmp195
    tmp198 = tmp196 / tmp197
    tmp199 = tmp195 - tmp198
    tmp200 = tl.broadcast_to(tmp199, [XBLOCK, RBLOCK])
    tmp202 = tl.sum(tmp200, 1)[:, None]
    tmp203 = tmp95 * tmp193
    tmp204 = tmp203 + tmp195
    tmp205 = tmp111 + tmp195
    tmp206 = tmp204 / tmp205
    tmp207 = tmp195 - tmp206
    tmp208 = tl.broadcast_to(tmp207, [XBLOCK, RBLOCK])
    tmp210 = tl.sum(tmp208, 1)[:, None]
    tmp211 = tmp136 * tmp193
    tmp212 = tmp211 + tmp195
    tmp213 = tmp152 + tmp195
    tmp214 = tmp212 / tmp213
    tmp215 = tmp195 - tmp214
    tmp216 = tl.broadcast_to(tmp215, [XBLOCK, RBLOCK])
    tmp218 = tl.sum(tmp216, 1)[:, None]
    tmp219 = tmp176 * tmp193
    tmp220 = tmp219 + tmp195
    tmp221 = tmp192 + tmp195
    tmp222 = tmp220 / tmp221
    tmp223 = tmp195 - tmp222
    tmp224 = tl.broadcast_to(tmp223, [XBLOCK, RBLOCK])
    tmp226 = tl.sum(tmp224, 1)[:, None]
    tmp227 = 4.0
    tmp228 = tmp202 / tmp227
    tmp229 = 0.0
    tmp230 = tmp228 + tmp229
    tmp231 = tmp210 / tmp227
    tmp232 = tmp230 + tmp231
    tmp233 = tmp218 / tmp227
    tmp234 = tmp232 + tmp233
    tmp235 = tmp226 / tmp227
    tmp236 = tmp234 + tmp235
    tmp237 = 0.25
    tmp238 = tmp236 * tmp237
    tmp239 = tmp238 / tmp195
    tmp240 = tmp239 * tmp195
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp240, None)


def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4), (4, 1))
    assert_size_stride(arg1_1, (4, 4), (4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused__softmax_0[grid(16)](arg0_1, buf0, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        del arg0_1
        buf1 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused__softmax_1[grid(16)](buf0, buf1, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        del buf0
        buf10 = empty_strided_cuda((), (), torch.float32)
        buf14 = buf10
        del buf10
        triton_per_fused__to_copy_add_div_mean_mul_ne_pow_rsub_sum_view_2[grid
            (1)](buf14, buf1, arg1_1, 1, 4, XBLOCK=1, num_warps=2, num_stages=1
            )
        del arg1_1
        del buf1
    return buf14,


def reduce_loss(loss, reduction):
    """Reduce loss as specified.

    Args:
        loss (Tensor): Elementwise loss tensor.
        reduction (str): Options are "none", "mean" and "sum".

    Return:
        Tensor: Reduced loss tensor.
    """
    reduction_enum = F._Reduction.get_enum(reduction)
    if reduction_enum == 0:
        return loss
    elif reduction_enum == 1:
        return loss.mean()
    elif reduction_enum == 2:
        return loss.sum()


def weight_reduce_loss(loss, weight=None, reduction='mean', avg_factor=None):
    """Apply element-wise weight and reduce loss.

    Args:
        loss (Tensor): Element-wise loss.
        weight (Tensor): Element-wise weights.
        reduction (str): Same as built-in losses of PyTorch.
        avg_factor (float): Avarage factor when computing the mean of losses.

    Returns:
        Tensor: Processed loss values.
    """
    if weight is not None:
        assert weight.dim() == loss.dim()
        if weight.dim() > 1:
            assert weight.size(1) == 1 or weight.size(1) == loss.size(1)
        loss = loss * weight
    if avg_factor is None:
        loss = reduce_loss(loss, reduction)
    elif reduction == 'mean':
        loss = loss.sum() / avg_factor
    elif reduction != 'none':
        raise ValueError('avg_factor can not be used with reduction="sum"')
    return loss


def get_class_weight(class_weight):
    """Get class weight for loss function.

    Args:
        class_weight (list[float] | str | None): If class_weight is a str,
            take it as a file name and read from it.
    """
    if isinstance(class_weight, str):
        if class_weight.endswith('.npy'):
            class_weight = np.load(class_weight)
        else:
            class_weight = mmcv.load(class_weight)
    return class_weight


def weighted_loss(loss_func):
    """Create a weighted version of a given loss function.

    To use this decorator, the loss function must have the signature like
    `loss_func(pred, target, **kwargs)`. The function only needs to compute
    element-wise loss without any reduction. This decorator will add weight
    and reduction arguments to the function. The decorated function will have
    the signature like `loss_func(pred, target, weight=None, reduction='mean',
    avg_factor=None, **kwargs)`.

    :Example:

    >>> import torch
    >>> @weighted_loss
    >>> def l1_loss(pred, target):
    >>>     return (pred - target).abs()

    >>> pred = torch.Tensor([0, 2, 3])
    >>> target = torch.Tensor([1, 1, 1])
    >>> weight = torch.Tensor([1, 0, 1])

    >>> l1_loss(pred, target)
    tensor(1.3333)
    >>> l1_loss(pred, target, weight)
    tensor(1.)
    >>> l1_loss(pred, target, reduction='none')
    tensor([1., 1., 2.])
    >>> l1_loss(pred, target, weight, avg_factor=2)
    tensor(1.5000)
    """

    @functools.wraps(loss_func)
    def wrapper(pred, target, weight=None, reduction='mean', avg_factor=
        None, **kwargs):
        loss = loss_func(pred, target, **kwargs)
        loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
        return loss
    return wrapper


@weighted_loss
def binary_dice_loss(pred, target, valid_mask, smooth=1, exponent=2, **kwards):
    assert pred.shape[0] == target.shape[0]
    pred = pred.reshape(pred.shape[0], -1)
    target = target.reshape(target.shape[0], -1)
    valid_mask = valid_mask.reshape(valid_mask.shape[0], -1)
    num = torch.sum(torch.mul(pred, target) * valid_mask, dim=1) * 2 + smooth
    den = torch.sum(pred.pow(exponent) + target.pow(exponent), dim=1) + smooth
    return 1 - num / den


@weighted_loss
def dice_loss(pred, target, valid_mask, smooth=1, exponent=2, class_weight=
    None, ignore_index=255):
    assert pred.shape[0] == target.shape[0]
    total_loss = 0
    num_classes = pred.shape[1]
    for i in range(num_classes):
        if i != ignore_index:
            dice_loss = binary_dice_loss(pred[:, i], target[..., i],
                valid_mask=valid_mask, smooth=smooth, exponent=exponent)
            if class_weight is not None:
                dice_loss *= class_weight[i]
            total_loss += dice_loss
    return total_loss / num_classes


class DiceLossNew(nn.Module):
    """DiceLoss.

    This loss is proposed in `V-Net: Fully Convolutional Neural Networks for
    Volumetric Medical Image Segmentation <https://arxiv.org/abs/1606.04797>`_.

    Args:
        loss_type (str, optional): Binary or multi-class loss.
            Default: 'multi_class'. Options are "binary" and "multi_class".
        smooth (float): A float number to smooth loss, and avoid NaN error.
            Default: 1
        exponent (float): An float number to calculate denominator
            value: \\sum{x^exponent} + \\sum{y^exponent}. Default: 2.
        reduction (str, optional): The method used to reduce the loss. Options
            are "none", "mean" and "sum". This parameter only works when
            per_image is True. Default: 'mean'.
        class_weight (list[float] | str, optional): Weight of each class. If in
            str format, read them from a file. Defaults to None.
        loss_weight (float, optional): Weight of the loss. Default to 1.0.
        ignore_index (int | None): The label index to be ignored. Default: 255.
        loss_name (str, optional): Name of the loss item. If you want this loss
            item to be included into the backward graph, `loss_` must be the
            prefix of the name. Defaults to 'loss_dice'.
    """

    def __init__(self, smooth=1, exponent=2, reduction='mean', class_weight
        =None, loss_weight=1.0, ignore_index=255, loss_name='loss_dice', **
        kwards):
        super(DiceLossNew, self).__init__()
        self.smooth = smooth
        self.exponent = exponent
        self.reduction = reduction
        self.class_weight = get_class_weight(class_weight)
        self.loss_weight = loss_weight
        self.ignore_index = ignore_index
        self._loss_name = loss_name

    @property
    def loss_name(self):
        """Loss Name.

        This function must be implemented and will return the name of this
        loss function. This name will be used to combine different loss items
        by simple sum operation. In addition, if you want this loss item to be
        included into the backward graph, `loss_` must be the prefix of the
        name.
        Returns:
            str: The name of this loss item.
        """
        return self._loss_name

    def forward(self, input_0, input_1):
        arg0_1 = input_0
        arg1_1 = input_1
        output = call([arg0_1, arg1_1])
        return output[0]


Measured latency_ms: 0.039902