mirror of
https://github.com/NVIDIA/cutlass.git
synced 2026-07-11 17:48:57 +00:00
64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
# 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_([])
|