import os
from pathlib import Path
import functools
import pandas as pd
import torch
import torch.nn.functional as F
from aiter import hipb_create_extension, hipb_mm, getHipblasltKernelName
from aiter import rocb_create_extension, rocb_mm
from aiter import logger, dtypes
from aiter.jit.utils.torch_guard import torch_compile_guard
from typing import Optional

extensions_created = False
@torch_compile_guard()
def scale_mm(
        inp: torch.Tensor,
        weights: torch.Tensor,
        bias: Optional[torch.Tensor] = None,
        otype: Optional[torch.dtype] = None,
        scale_a: Optional[torch.Tensor] = None,
        scale_b: Optional[torch.Tensor] = None,
        scale_c: Optional[torch.Tensor] = None,
        scale_type: Optional[int] = None,
)-> torch.Tensor:
    # scale_type=0, scalar  scale
    # scale_type=1, channel scale
    # scale_type=2, block   scale
    global extensions_created
    if otype is None:
        otype = inp.dtype
    if extensions_created == False:
        hipb_create_extension()
        extensions_created = True
    if inp.dim() >= 3:
        assert(False, "not support 3dim input")

    inp_view = inp
    return hipb_mm(inp_view, weights.t(), -1, bias, otype, scale_a, scale_b, scale_c, scale_type)