#!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
import os
from torch.utils.cpp_extension import load_inline

# Set up environment variables for HIP compilation
os.environ["CXX"] = "clang++"
os.environ["PYTORCH_ROCM_ARCH"] = "gfx942"

# HIP kernel for SiLU matmul
cuda_src = r"""
#include <iostream>
#include <hip/hip_runtime.h>
#include <hip/amd_detail/amd_hip_bf16.h>

// Optimized HIP kernel for SiLU activation function followed by matrix multiplication
template <typename T>
__global__ void silu_matmul_kernel(
    const T* __restrict__ input,     // [total_tokens, hidden_dim]
    const T* __restrict__ weights,   // [hidden_dim, output_dim]
    T* __restrict__ output,          // [total_tokens, output_dim]
    const int total_tokens,
    const int hidden_dim,
    const int output_dim
) {
    // Use shared memory to reduce global memory access
    extern __shared__ float shared_mem[];
    
    const int TILE_DIM = 16;
    const int row = blockIdx.x * TILE_DIM + threadIdx.x;
    const int col = blockIdx.y * TILE_DIM + threadIdx.y;
    
    // Each thread computes one element of the output matrix
    float sum = 0.0f;
    
    // Loop over tiles to handle large matrices
    for (int t = 0; t < (hidden_dim + TILE_DIM - 1) / TILE_DIM; ++t) {
        // Each thread loads one element from input and weight matrices
        if (row < total_tokens && t * TILE_DIM + threadIdx.y < hidden_dim) {
            shared_mem[threadIdx.x * TILE_DIM + threadIdx.y] = 
                static_cast<float>(input[row * hidden_dim + t * TILE_DIM + threadIdx.y]);
        } else {
            shared_mem[threadIdx.x * TILE_DIM + threadIdx.y] = 0.0f;
        }
        
        if (t * TILE_DIM + threadIdx.x < hidden_dim && col < output_dim) {
            shared_mem[TILE_DIM * TILE_DIM + threadIdx.x * TILE_DIM + threadIdx.y] = 
                static_cast<float>(weights[(t * TILE_DIM + threadIdx.x) * output_dim + col]);
        } else {
            shared_mem[TILE_DIM * TILE_DIM + threadIdx.x * TILE_DIM + threadIdx.y] = 0.0f;
        }
        
        __syncthreads();
        
        // Compute partial product for this tile
        #pragma unroll
        for (int k = 0; k < TILE_DIM; ++k) {
            sum += shared_mem[threadIdx.x * TILE_DIM + k] * 
                   shared_mem[TILE_DIM * TILE_DIM + k * TILE_DIM + threadIdx.y];
        }
        
        __syncthreads();
    }
    
    // Apply SiLU activation (x * sigmoid(x)) and write to output
    if (row < total_tokens && col < output_dim) {
        // SiLU: x * sigmoid(x)
        const float sigmoid_x = 1.0f / (1.0f + expf(-sum));
        output[row * output_dim + col] = static_cast<T>(sum * sigmoid_x);
    }
}

// Wrapper function for the SiLU matmul operation
at::Tensor silu_matmul(
    at::Tensor input,
    at::Tensor weights
) {
    // Get tensor dimensions
    const int total_tokens = input.size(0);
    const int hidden_dim = input.size(1);
    const int output_dim = weights.size(1);
    
    // Create output tensor
    at::Tensor output = at::empty({total_tokens, output_dim}, input.options());
    
    // Set up thread blocks and grid
    const int TILE_DIM = 16;
    dim3 threads(TILE_DIM, TILE_DIM);
    dim3 grid(
        (total_tokens + TILE_DIM - 1) / TILE_DIM,
        (output_dim + TILE_DIM - 1) / TILE_DIM
    );
    
    // Calculate shared memory size (2 tiles)
    const int sharedMemSize = 2 * TILE_DIM * TILE_DIM * sizeof(float);
    
    // Launch kernel based on data type
    if (input.scalar_type() == at::ScalarType::Float) {
        silu_matmul_kernel<float><<<grid, threads, sharedMemSize>>>(
            input.data_ptr<float>(),
            weights.data_ptr<float>(),
            output.data_ptr<float>(),
            total_tokens,
            hidden_dim,
            output_dim
        );
    } else if (input.scalar_type() == at::ScalarType::Half) {
        silu_matmul_kernel<at::Half><<<grid, threads, sharedMemSize>>>(
            reinterpret_cast<at::Half*>(input.data_ptr()),
            reinterpret_cast<at::Half*>(weights.data_ptr()),
            reinterpret_cast<at::Half*>(output.data_ptr()),
            total_tokens,
            hidden_dim,
            output_dim
        );
    } else if (input.scalar_type() == at::ScalarType::BFloat16) {
        silu_matmul_kernel<__hip_bfloat16><<<grid, threads, sharedMemSize>>>(
            reinterpret_cast<__hip_bfloat16*>(input.data_ptr()),
            reinterpret_cast<__hip_bfloat16*>(weights.data_ptr()),
            reinterpret_cast<__hip_bfloat16*>(output.data_ptr()),
            total_tokens,
            hidden_dim,
            output_dim
        );
    } else {
        AT_ERROR("Unsupported data type for silu_matmul");
    }
    
    // Check for CUDA errors
    cudaError_t error = cudaGetLastError();
    if (error != cudaSuccess) {
        AT_ERROR("CUDA error: ", cudaGetErrorString(error));
    }
    
    return output;
}
"""

cpp_src = r"""
at::Tensor silu_matmul(at::Tensor input, at::Tensor weights);
"""

# Load the custom kernel
silu_matmul_module = load_inline(
    name="silu_matmul",
    cpp_sources=[cpp_src],
    cuda_sources=[cuda_src],
    functions=["silu_matmul"],
    verbose=True,
    extra_cuda_cflags=[
        "-Ofast",
        "--offload-arch=gfx942",
        "-std=c++17",
        "-ffp-contract=fast",
        "-lhip_hcc",
        "-mcumode",
    ],
)

def hip_silu_matmul(x, weight):
    """
    HIP implementation of SiLU matmul
    Args:
        x: Input tensor [batch_size, hidden_dim]
        weight: Weight tensor [hidden_dim, output_dim]
    Returns:
        Output tensor after SiLU activation [batch_size, output_dim]
    """
    return silu_matmul_module.silu_matmul(x, weight)


class Fused_MOE(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        
        self.num_routed_experts = config["n_routed_experts"]
        self.num_experts_per_token = config["n_experts_per_token"]
        
        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:
        orig_shape = x.shape
        batch_size, seq_len, hidden_dim = x.shape
        total_tokens = batch_size * seq_len
        
        # Reshape x for matmul
        x_flat = x.reshape(total_tokens, hidden_dim)
        
        # Use HIP implementation for SiLU matmul
        shared_gate_activation = hip_silu_matmul(
            x_flat, 
            self.all_weights["shared_gate"]
        )
        
        # Regular matmul for shared_up
        shared_up = x_flat @ self.all_weights["shared_up"]
        
        router_logits = x_flat @ self.all_weights["router"]
        
        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
        )
        
        shared_output = (shared_gate_activation * shared_up) @ self.all_weights["shared_down"]
        
        flat_topk_indices = topk_indices.reshape(-1)
        flat_topk_probs = topk_probs.reshape(-1, 1)
        
        routed_output = torch.zeros_like(x_flat)
        
        sorted_idx = flat_topk_indices.argsort()
        token_idx = sorted_idx // self.num_experts_per_token
        
        expert_counts = flat_topk_indices.bincount()
        expert_cumsum = expert_counts.cumsum(0).cpu().numpy()
        
        for expert_id in range(self.num_routed_experts):
            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
            
            if start_idx == end_idx:
                continue
                
            current_tokens = token_idx[start_idx:end_idx]
            
            expert_input = x_flat[current_tokens]
            
            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]
            
            # Use HIP implementation for expert SiLU matmul
            expert_gate_activation = hip_silu_matmul(expert_input, gate_weight)
            expert_up = expert_input @ up_weight
            
            expert_output = (expert_gate_activation * expert_up) @ down_weight
            expert_output.mul_(flat_topk_probs[sorted_idx[start_idx:end_idx]])
            
            expanded_indices = current_tokens.unsqueeze(1).expand(-1, expert_output.size(1))
            routed_output.scatter_add_(0, expanded_indices, expert_output)
            
        return routed_output.reshape(orig_shape) + shared_output.reshape(orig_shape)
        
    def load_weights(self, weights):
        self.all_weights["router"] = weights["router.weight"].t()
        
        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"]
        
        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
    
    fused_moe_instance = Fused_MOE(config)
    
    fused_moe_instance.load_weights(weights)
    
    output = fused_moe_instance(input_tensor) 
    
    return output
python submission.py
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/silu_matmul/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/silu_matmul/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/silu_matmul/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/silu_matmul/hip.hip [ok]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 3
[1/3] clang++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=silu_matmul -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -c /home/runner/.cache/torch_extensions/py310_cpu/silu_matmul/main.cpp -o main.o -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC
[2/3] /opt/rocm-6.3.1/bin/hipcc  -DWITH_HIP -DTORCH_EXTENSION_NAME=silu_matmul -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -isystem /usr/local/lib/python3.10/dist-packages/torch/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/torch/csrc/api/include -isystem /usr/local/lib/python3.10/dist-packages/torch/include/THH -isystem /opt/rocm-6.3.1/include -isystem /usr/include/python3.10 -fPIC -std=c++17 -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC -DCUDA_HAS_FP16=1 -D__HIP_NO_HALF_OPERATORS__=1 -D__HIP_NO_HALF_CONVERSIONS__=1 -DHIP_ENABLE_WARP_SYNC_BUILTINS=1 --offload-arch=gfx942 -fno-gpu-rdc -Ofast --offload-arch=gfx942 -std=c++17 -ffp-contract=fast -lhip_hcc -mcumode -c /home/runner/.cache/torch_extensions/py310_cpu/silu_matmul/hip.hip -o hip.cuda.o 
clang++: warning: -lhip_hcc: 'linker' input unused [-Wunused-command-line-argument]
[3/3] clang++ main.o hip.cuda.o -shared -L/usr/local/lib/python3.10/dist-packages/torch/lib -lc10 -lc10_hip -ltorch_cpu -ltorch_hip -ltorch -ltorch_python -L/opt/rocm-6.3.1/lib -lamdhip64 -o silu_matmul.so

python eval.py benchmark /tmp/tmpwz6ogjx2
[92mSuccessfully preprocessed all matching files.[0m

/home/runner/.cache/torch_extensions/py310_cpu/silu_matmul/main.cpp -> /home/runner/.cache/torch_extensions/py310_cpu/silu_matmul/main.cpp [skipped, no changes]
/home/runner/.cache/torch_extensions/py310_cpu/silu_matmul/cuda.cu -> /home/runner/.cache/torch_extensions/py310_cpu/silu_matmul/hip.hip [skipped, already hipified]
Total number of unsupported CUDA function calls: 0


Total number of replaced kernel launches: 3
ninja: no work to do.
