# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Use of this software is governed by the terms and conditions of the
# NVIDIA End User License Agreement (EULA), available at:
# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
#
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.

"""
This module provides MLIR GPU Dialect helper functions
"""

from .._mlir import ir
from .._mlir.dialects import gpu, arith, scf
from .._mlir.extras import types as _T

from ..base_dsl.common import *

# =============================================================================
# GPU Dialect Helper functions
# =============================================================================


def create_async_token() -> ir.Value:
    token_ty = gpu.AsyncTokenType.get()
    token = gpu.wait(token_ty, [])
    return token


def printf(fmt: str, *args: ir.Value, threadNumber: int = -1) -> None:
    """Generate gpu.printf OP predicated on threadNumber"""
    type_formats = []
    for arg in args:
        ty_format = None
        if isinstance(arg.type, ir.IndexType):
            ty_format = "%llu"
        if isinstance(arg.type, ir.IntegerType):
            width = ir.IntegerType(arg.type).width
            if width == 64:
                ty_format = "%llu"
            elif width == 32:
                ty_format = "%d"
            elif width == 1:
                ty_format = "%i"
        if isinstance(arg.type, ir.F32Type):
            ty_format = "%f"
        if ty_format is None:
            raise DSLNotImplemented(arg.type)
        type_formats.append(ty_format)
    if threadNumber == -1:
        gpu.printf(fmt.format(*type_formats) + "\n", args)
    if threadNumber != -1:
        tidx = gpu.thread_id(gpu.Dimension.x)
        predicate = arith.cmpi(
            arith.CmpIPredicate.eq, tidx, arith.constant(_T.index(), threadNumber)
        )
        if_op = scf.IfOp(predicate)
        with ir.InsertionPoint(if_op.then_block):
            gpu.printf(fmt.format(*type_formats) + "\n", args)
            scf.yield_([])