You write custom CUDA kernels to replace pytorch operators in given architecture to get speedups. You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.

Here's an example to show you the syntax of inline embedding custom CUDA operators in PyTorch. The example given architecture is a simple addition:

python
import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
def init(self) -> None:
super().init()

def forward(self, a, b):
    return a + b
def get_inputs():
a = torch.randn(1, 128).cuda()
b = torch.randn(1, 128).cuda()
return [a, b]

def get_init_inputs():
return []



The example new architecture with a custom CUDA kernel looks like this:

python
import torch
from torch.utils.cpp_extension import load_inline

add_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

global void add_kernel(const float* a, const float* b, float* out, int size) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < size) {
out[idx] = a[idx] + b[idx];
}
}

torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b) {
auto out = torch::empty_like(a);
int size = a.numel();
const int block_size = 256;
int num_blocks = (size + block_size - 1) / block_size;
add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);
return out;
}
"""

add_cpp_source = """
torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b);
"""

Compile the inline CUDA code
add = load_inline(
name="add",
cpp_sources=add_cpp_source,
cuda_sources=add_source,
functions=["add_cuda"],
verbose=True
)

class ModelNew(torch.nn.Module):
def init(self):
super(ModelNew, self).init()
self.add = add

def forward(self, a, b):
    return self.add.add_cuda(a, b)


---

Now, you are given the following PyTorch architecture to accelerate. The model computes a combined loss, which is a weighted sum of Dice Loss and Binary Cross-Entropy with Logits Loss (BCEWithLogitsLoss). This is a very common pattern in segmentation tasks. The baseline implementation uses separate, highly optimized PyTorch functions for each part.

python
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    PyTorch基准实现：Dice Loss + BCEWithLogitsLoss
    """
    def __init__(self, dice_weight=1.0, bce_weight=1.0):
        super(Model, self).__init__()
        self.dice_weight = dice_weight
        self.bce_weight = bce_weight
        # BCEWithLogitsLoss内部融合了Sigmoid和BCE，更稳定
        self.bce_loss_fn = nn.BCEWithLogitsLoss(reduction='mean')

    def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """
        Computes the combined Dice and BCE loss.

        Args:
            logits (torch.Tensor): Raw prediction tensor of shape [N, C, H, W].
            target (torch.Tensor): Target tensor of same shape as logits.

        Returns:
            torch.Tensor: Combined loss value (scalar).
        """
        # --- 第一步：计算Dice Loss ---
        pred = torch.sigmoid(logits)
        pred_flat = pred.view(-1)
        target_flat = target.view(-1)
        
        intersection = (pred_flat * target_flat).sum()
        pred_sum = pred_flat.sum()
        target_sum = target_flat.sum()
        
        epsilon = 1e-6
        dice_score = (2.0 * intersection) / (pred_sum + target_sum + epsilon)
        loss_dice = 1.0 - dice_score
        
        # --- 第二步：计算BCEWithLogitsLoss ---
        # PyTorch的BCEWithLogitsLoss已经非常高效且稳定
        loss_bce = self.bce_loss_fn(logits, target)
        
        # --- 第三步：组合损失 ---
        loss_total = self.dice_weight * loss_dice + self.bce_weight * loss_bce
        
        return loss_total

batch_size = 32
height, width = 256, 256
channels = 1

def get_inputs():
    logits = torch.randn(batch_size, channels, height, width)
    target = torch.randint(0, 2, (batch_size, channels, height, width), dtype=torch.float32)
    return [logits, target]

def get_init_inputs():
    return [] # No special initialization inputs needed



Your task is to generate the `ModelNew` architecture with a single, highly optimized custom CUDA kernel that fuses the computation of both the Dice Loss and the BCEWithLogitsLoss.

**CRITICAL REQUIREMENTS:**

1.  **Performance Optimization & Fusion:**
    *   **Operator Fusion:** The entire calculation for both losses must be performed within a **single CUDA kernel**. The kernel should take `logits` and `target` as input and compute all necessary intermediate sums.
    *   The kernel should use a **Warp-level reduction** strategy for optimal performance. Warps should collaboratively iterate over the total number of elements.
    *   **Data Reuse:** The Sigmoid of `logits` should be computed once per element and then reused for both the Dice and BCE parts of the calculation. This avoids redundant computation and memory traffic.

2.  **Kernel Logic:**
    *   The kernel should launch a grid of blocks where the total number of threads is sufficient to cover all elements in the input tensors (`logits.numel()`).
    *   Each thread should process one element. Inside the loop, it should compute `p = sigmoid(logits[i])`.
    *   It should then accumulate the necessary values for the final loss calculation: `intersection += p * target[i]`, `pred_sum += p`, `target_sum += target[i]`, and `bce_sum += BCE_loss(p, target[i])`.
    *   Use `__shfl_down_sync` for efficient warp-level reduction to compute the final four sums: `intersection`, `pred_sum`, `target_sum`, and `bce_sum`.
    *   The host-side C++ wrapper should then compute the final `loss_dice`, `loss_bce`, and `loss_total` from these four sums. Use the numerically stable formula for BCE: `-[t * log(p) + (1-t) * log(1-p)]`.

3.  **Code Structure:** Follow the exact structure of the provided example, including `load_inline`, the CUDA source string, the C++ wrapper source string, and the `ModelNew` class. The `get_init_inputs` function must return `[]` to match the baseline.

4.  **Compilation Flags:** Use `-O3` and `--use_fast_math` for maximum performance, as this is a common practice for high-throughput kernels like this. Avoid hardcoding compute capabilities to ensure portability.
