#!POPCORN leaderboard amd-mixture-of-experts
#!POPCORN gpus MI300

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, Optional
from task import input_t, output_t


class HyperFusedMoE(nn.Module):
    """Ultra-fused implementation with extreme kernel fusion and memory optimizations."""
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        
        # Cache configuration values for faster access
        self.num_routed_experts = config["n_routed_experts"]
        self.num_experts_per_token = config["n_experts_per_token"]
        
        # Pre-allocate weights - will be populated in custom_kernel
        self.all_weights = {
            "router": None,
            "experts_gate": [None] * self.num_routed_experts,
            "experts_up": [None] * self.num_routed_experts,
            "experts_down": [None] * self.num_routed_experts,
            "shared_gate": None,
            "shared_up": None,
            "shared_down": None
        }

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Preserve original shape for final reshape
        orig_shape = x.shape
        batch_size, seq_len, hidden_dim = x.shape
        total_tokens = batch_size * seq_len
        
        # =====================================================
        # STAGE 1: Compute all linear projections in one batch
        # =====================================================
        # This reduces memory bandwidth by doing multiple matrix multiplications
        # at once before moving to the next stage
        
        # Compute shared expert projections and router logits in a fused manner
        # These are all matrix multiplications with the same input (x)
        shared_gate_activation = F.silu(x @ self.all_weights["shared_gate"])
        shared_up = x @ self.all_weights["shared_up"]
        router_logits = x @ self.all_weights["router"]
        
        # Compute router probabilities and select top-k experts
        # Use in-place operations where possible
        router_probs = F.softmax(router_logits, dim=-1)
        topk_probs, topk_indices = torch.topk(
            router_probs, k=self.num_experts_per_token, dim=-1, sorted=False
        )
        
        # =====================================================
        # STAGE 2: Compute shared expert output (fused operations)
        # =====================================================
        # Complete shared expert computation while router results are fresh in cache
        shared_output = (shared_gate_activation * shared_up) @ self.all_weights["shared_down"]
        
        # =====================================================
        # STAGE 3: Prepare for expert routing with minimal reshape operations
        # =====================================================
        # Flatten input for expert computation
        x_flat = x.reshape(total_tokens, hidden_dim)
        
        # Flatten expert indices and probabilities
        flat_topk_indices = topk_indices.reshape(-1)
        flat_topk_probs = topk_probs.reshape(-1, 1)
        
        # =====================================================
        # STAGE 4: Hyper-fused expert routing and computation
        # =====================================================
        # Allocate output buffer
        routed_output = torch.zeros_like(x_flat)
        
        # Sort indices for better memory locality during expert processing
        sorted_idx = flat_topk_indices.argsort()
        token_idx = sorted_idx // self.num_experts_per_token  # Original token indices
        
        # Count tokens per expert for batch processing
        expert_counts = flat_topk_indices.bincount()
        expert_cumsum = expert_counts.cumsum(0).cpu().numpy()
        
        # Process each expert in a fully fused manner
        for expert_id in range(self.num_routed_experts):
            # Get token range for this expert
            start_idx = 0 if expert_id == 0 else expert_cumsum[expert_id - 1]
            end_idx = expert_cumsum[expert_id] if expert_id < len(expert_cumsum) else 0
            
            # Skip if no tokens for this expert
            if start_idx == end_idx:
                continue
                
            # Get the actual token indices for this expert batch
            current_tokens = token_idx[start_idx:end_idx]
            
            # Get input tokens in a single operation
            expert_input = x_flat[current_tokens]
            
            # Get weights for this expert
            gate_weight = self.all_weights["experts_gate"][expert_id]
            up_weight = self.all_weights["experts_up"][expert_id]
            down_weight = self.all_weights["experts_down"][expert_id]
            
            # Fused expert computation - compute everything in one go
            # 1. Compute gate activation
            # 2. Compute up projection
            # 3. Element-wise multiply
            # 4. Compute down projection
            # 5. Scale by router probabilities
            expert_output = F.silu(expert_input @ gate_weight) * (expert_input @ up_weight) @ down_weight
            expert_output.mul_(flat_topk_probs[sorted_idx[start_idx:end_idx]])
            
            # Efficient scatter_add to accumulate results
            # Use expanded indices to efficiently handle the hidden dimension
            expanded_indices = current_tokens.unsqueeze(1).expand(-1, expert_output.size(1))
            routed_output.scatter_add_(0, expanded_indices, expert_output)
            
        # =====================================================
        # STAGE 5: Combine shared and routed outputs
        # =====================================================
        # Reshape routed output to original dimensions and add shared output
        return routed_output.reshape(orig_shape) + shared_output
        
    def load_weights(self, weights):
        """Load all weights from the weight dictionary."""
        # Router weights
        self.all_weights["router"] = weights["router.weight"].t()
        
        # Expert weights
        for i in range(self.num_routed_experts):
            self.all_weights["experts_gate"][i] = weights[f"experts.{i}.0.weight"]
            self.all_weights["experts_up"][i] = weights[f"experts.{i}.1.weight"]
            self.all_weights["experts_down"][i] = weights[f"experts.{i}.2.weight"]
        
        # Shared expert weights
        self.all_weights["shared_gate"] = weights["shared_experts.0.weight"]
        self.all_weights["shared_up"] = weights["shared_experts.1.weight"]
        self.all_weights["shared_down"] = weights["shared_experts.2.weight"]


def custom_kernel(data: input_t) -> output_t:
    input_tensor, weights, config = data
    
    # Create hyperfused MoE
    hyperfused_moe = HyperFusedMoE(config)
    
    # Load all weights at once
    hyperfused_moe.load_weights(weights)
    
    # Forward pass
    output = hyperfused_moe(input_tensor)
    
    return output
python eval.py benchmark /tmp/tmp4c4manyd