Release v4.0.0 (#2294)
This commit is contained in:
Kihiro Bando
@@ -0,0 +1,221 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 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/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 a class that compiles generated IR using MLIR's PassManager
|
||||
and executes it using MLIR's ExecutionEngine.
|
||||
|
||||
"""
|
||||
|
||||
from typing import Sequence, Optional, Tuple
|
||||
import os
|
||||
import sys
|
||||
import inspect
|
||||
from .common import DSLRuntimeError
|
||||
|
||||
_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
|
||||
sys.path.append(_SCRIPT_PATH)
|
||||
|
||||
from .._mlir import ir
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Compiler Class
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class CompilationError(RuntimeError):
|
||||
"""Custom error class for compilation failures"""
|
||||
|
||||
# Add ANSI color codes
|
||||
RED = "\033[91m"
|
||||
YELLOW = "\033[93m"
|
||||
BLUE = "\033[94m"
|
||||
GREEN = "\033[92m"
|
||||
BOLD = "\033[1m"
|
||||
RESET = "\033[0m"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
nvvm_error: Optional[str] = None,
|
||||
ir_context: Optional[str] = None,
|
||||
cuda_toolkit: Optional[str] = None,
|
||||
arch: Optional[str] = None,
|
||||
):
|
||||
self.nvvm_error = nvvm_error
|
||||
self.ir_context = ir_context
|
||||
self.cuda_toolkit = cuda_toolkit
|
||||
self.arch = arch
|
||||
# Call parent with formatted error to avoid showing class name
|
||||
super().__init__("") # Empty string to avoid class name
|
||||
# Store formatted error for str() representation
|
||||
self._formatted_error = self._format_error()
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""Override string representation to avoid showing class name"""
|
||||
return self._formatted_error
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""Override repr representation to avoid showing class name"""
|
||||
return self._formatted_error
|
||||
|
||||
def _format_error(self) -> str:
|
||||
if not self.nvvm_error:
|
||||
return str(self.args[0])
|
||||
|
||||
return f"""NVVM Compilation Error:
|
||||
----------------------
|
||||
|
||||
{self.BLUE}⚙️ Current Settings:{self.RESET}
|
||||
{self.BOLD}- CUDA Toolkit Path: {self.cuda_toolkit or "Not Set"}
|
||||
- Target Architecture: {self.arch}{self.RESET}
|
||||
|
||||
IR Context (truncated):
|
||||
{self.ir_context}
|
||||
|
||||
{self.YELLOW}💡 Possible Solutions:{self.RESET}
|
||||
{self.GREEN}1. Check if CUDA_TOOLKIT_PATH is set correctly
|
||||
2. Verify target architecture ({self.arch}) is supported by your CUDA toolkit
|
||||
3. Make sure CUDA toolkit version matches the target architecture requirements{self.RESET}"""
|
||||
|
||||
|
||||
class Compiler:
|
||||
"""Compiler class for compiling and building MLIR modules."""
|
||||
|
||||
def __init__(self, passmanager, execution_engine):
|
||||
self.passmanager = passmanager
|
||||
self.execution_engine = execution_engine
|
||||
|
||||
def __call__(self, module):
|
||||
"""Convenience application method."""
|
||||
self.compile(module)
|
||||
|
||||
def _process_error(self, error_msg: str) -> Tuple[Optional[str], Optional[str]]:
|
||||
"""Process error message to extract NVVM error and IR context"""
|
||||
nvvm_error = None
|
||||
ir_msg = ""
|
||||
|
||||
if "NVVM_ERROR" in error_msg:
|
||||
# Extract the specific NVVM error
|
||||
nvvm_error = (
|
||||
error_msg.split("libNVVM extra log:")[1].strip()
|
||||
if "libNVVM extra log:" in error_msg
|
||||
else error_msg
|
||||
)
|
||||
|
||||
# Extract IR context
|
||||
if "see current operation:" in error_msg:
|
||||
# Get the IR section
|
||||
ir_section = error_msg.split("see current operation:")[1].strip()
|
||||
# Remove duplicate IR section
|
||||
ir_section = ir_section.split("error: unknown: Failed translating")[
|
||||
0
|
||||
].strip()
|
||||
|
||||
# Get first few lines and last few lines of the IR
|
||||
ir_lines = ir_section.split("\n")
|
||||
if len(ir_lines) > 10:
|
||||
ir_msg = "\n".join(ir_lines[:5] + [" ..."] + ir_lines[-5:])
|
||||
else:
|
||||
ir_msg = ir_section
|
||||
|
||||
return nvvm_error, ir_msg
|
||||
|
||||
def compile(
|
||||
self,
|
||||
module,
|
||||
pipeline: str,
|
||||
cuda_toolkit: str = "",
|
||||
arch: str = "",
|
||||
enable_verifier=False,
|
||||
):
|
||||
"""Compiles the module by invoking the pipeline."""
|
||||
try:
|
||||
pm = self.passmanager.PassManager.parse(pipeline)
|
||||
pm.enable_verifier(enable_verifier)
|
||||
pm.run(module.operation)
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
nvvm_error, ir_msg = self._process_error(error_msg)
|
||||
|
||||
if nvvm_error:
|
||||
raise CompilationError(
|
||||
error_msg,
|
||||
nvvm_error=nvvm_error,
|
||||
ir_context=ir_msg,
|
||||
cuda_toolkit=cuda_toolkit,
|
||||
arch=arch,
|
||||
) from e
|
||||
raise e
|
||||
|
||||
def jit(self, module, opt_level: int = 2, shared_libs: Sequence[str] = ()):
|
||||
"""Wraps the module in a JIT execution engine."""
|
||||
return self.execution_engine.ExecutionEngine(
|
||||
module, opt_level=opt_level, shared_libs=shared_libs
|
||||
)
|
||||
|
||||
def compile_and_jit(
|
||||
self,
|
||||
module,
|
||||
pipeline: str,
|
||||
shared_libs: Sequence[str] = (),
|
||||
opt_level: int = 2,
|
||||
cuda_toolkit: str = "",
|
||||
arch: str = "",
|
||||
):
|
||||
"""Compiles and jits the module."""
|
||||
self.compile(
|
||||
module,
|
||||
pipeline,
|
||||
cuda_toolkit,
|
||||
arch,
|
||||
)
|
||||
return self.jit(module, opt_level, shared_libs)
|
||||
|
||||
|
||||
def compile(func, *args, **kwargs):
|
||||
if func is None:
|
||||
raise DSLRuntimeError("Function is not set or invalid.")
|
||||
|
||||
if not callable(func):
|
||||
raise DSLRuntimeError("Object is not callable.")
|
||||
|
||||
kwargs["compile_only"] = True
|
||||
kwargs["no_cache"] = True
|
||||
|
||||
if inspect.isfunction(func):
|
||||
# regular function
|
||||
pass
|
||||
elif inspect.ismethod(func):
|
||||
# if it's a method, add the instance to the first argument
|
||||
args = [func.__self__] + list(args)
|
||||
func = func.__func__
|
||||
elif inspect.isclass(type(func)) and hasattr(func, "__call__"):
|
||||
# If it's a class instance, get the class's __call__ method
|
||||
args = [func] + list(args)
|
||||
# Get the actual function from the class definition
|
||||
func = func.__call__.__func__
|
||||
else:
|
||||
raise DSLRuntimeError(
|
||||
"Invalid function type, only function, method and module are supported, but got",
|
||||
func,
|
||||
)
|
||||
|
||||
# If it's a wrapped function created by jit decorator, get the original function
|
||||
if hasattr(func, "__wrapped__"):
|
||||
func = func.__wrapped__
|
||||
|
||||
if not hasattr(func, "_dsl_object"):
|
||||
raise DSLRuntimeError("Function is not decorated with jit decorator.")
|
||||
|
||||
fcn_ptr = func._dsl_object._preprocess_and_execute(func)
|
||||
return func._dsl_object._func(fcn_ptr, *args, **kwargs)
|
||||
Reference in New Issue
Block a user