Merge pull request #570 from InfiniTensor/issue/564

issue/564 - updated test info, env settings, and readme

Author: PanZezhong1725

README.md

@@ -50,6 +50,13 @@ git clone --recursive https://github.com/InfiniTensor/InfiniCore.git
 git submodule update --init --recursive
 ```
 
+配置`INFINI_ROOT` 和 `LD_LIBRARY_PATH` 环境变量。  
+默认`INFINI_ROOT`为`$HOME/.infini`，可以使用以下命令自动配置：
+
+```shell
+source scripts/set_env_linux.sh
+```
+
 如果你需要在本地开发九齿算子（即需要对九齿算子库进行修改），推荐单独克隆[九齿算子库](https://github.com/InfiniTensor/ntops)，并从本地安装：
 
 ```shell
@@ -148,10 +155,6 @@ python scripts/install.py [XMAKE_CONFIG_FLAGS]
    xmake build && xmake install
    ```
 
-3. 设置环境变量
-
-   按输出提示设置 `INFINI_ROOT` 和 `LD_LIBRARY_PATH` 环境变量。
-
 #### 2. 安装 C++ 库
 
 ```shell

scripts/set_env_linux.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+
+# Define the content to be added
+content='export INFINI_ROOT="$HOME/.infini"
+export LD_LIBRARY_PATH="$INFINI_ROOT/lib:$LD_LIBRARY_PATH"'
+
+# Check if bashrc file exists
+bashrc_file="$HOME/.bashrc"
+
+if [ ! -f "$bashrc_file" ]; then
+    echo "Creating $bashrc_file file"
+    touch "$bashrc_file"
+fi
+
+# Check if the content already exists
+if grep -q "export INFINI_ROOT=" "$bashrc_file"; then
+    echo "INFINI_ROOT configuration already exists in $bashrc_file"
+else
+    echo "Adding configuration to $bashrc_file"
+    echo "$content" >> "$bashrc_file"
+    echo "Configuration added successfully"
+fi
+
+# Reload bashrc
+echo "Reloading $bashrc_file"
+source "$bashrc_file"
+
+echo "Done! INFINI_ROOT and LD_LIBRARY_PATH have been set"

test/infinicore/framework/__init__.py

@@ -9,33 +9,49 @@
     profile_operation,
     rearrange_tensor,
     convert_infinicore_to_torch,
+    get_operator_help_info,
+    print_operator_testing_tips,
+)
+from .config import (
+    get_args,
+    get_hardware_args_group,
+    get_hardware_help_text,
+    get_supported_hardware_platforms,
+    get_test_devices,
 )
-from .config import get_test_devices, get_args
 from .devices import InfiniDeviceEnum, InfiniDeviceNames, torch_device_map
 from .datatypes import to_torch_dtype, to_infinicore_dtype
 from .runner import GenericTestRunner
 
 __all__ = [
-    "TensorSpec",
+    # Core types and classes
+    "BaseOperatorTest",
+    "GenericTestRunner",
+    "InfiniDeviceEnum",
+    "InfiniDeviceNames",
     "TensorInitializer",
+    "TensorSpec",
+    "TestCase",
     "TestConfig",
     "TestRunner",
-    "TestCase",
-    "BaseOperatorTest",
+    # Core functions
     "compare_results",
-    "create_test_comparator",
     "convert_infinicore_to_torch",
+    "create_test_comparator",
     "debug",
     "get_args",
+    "get_hardware_args_group",
+    "get_hardware_help_text",
+    "get_operator_help_info",
+    "get_supported_hardware_platforms",
     "get_test_devices",
     "get_tolerance",
     "infinicore_tensor_from_torch",
+    "print_operator_testing_tips",
     "profile_operation",
     "rearrange_tensor",
-    "InfiniDeviceEnum",
-    "InfiniDeviceNames",
-    "torch_device_map",
-    "to_torch_dtype",
+    # Utility functions
     "to_infinicore_dtype",
-    "GenericTestRunner",
+    "to_torch_dtype",
+    "torch_device_map",
 ]

test/infinicore/framework/config.py

@@ -1,78 +1,150 @@
 import argparse
 from .devices import InfiniDeviceEnum
 
+# hardware_info.py
+"""
+Shared hardware platform information for the InfiniCore testing framework
+"""
+
+
+def get_supported_hardware_platforms():
+    """
+    Get list of supported hardware platforms with descriptions.
+
+    Returns:
+        List of tuples (flag, description)
+    """
+    return [
+        ("--cpu", "Standard CPU execution"),
+        ("--nvidia", "NVIDIA GPUs with CUDA support"),
+        ("--cambricon", "Cambricon MLU accelerators (requires torch_mlu)"),
+        ("--ascend", "Huawei Ascend NPUs (requires torch_npu)"),
+        ("--iluvatar", "Iluvatar GPUs"),
+        ("--metax", "Metax GPUs"),
+        ("--moore", "Moore Threads GPUs (requires torch_musa)"),
+        ("--kunlun", "Kunlun XPUs (requires torch_xmlir)"),
+        ("--hygon", "Hygon DCUs"),
+    ]
+
+
+def get_hardware_help_text():
+    """
+    Get formatted help text for hardware platforms.
+
+    Returns:
+        str: Formatted help text for argument parsers
+    """
+    platforms = get_supported_hardware_platforms()
+    help_lines = ["Supported Hardware Platforms:"]
+
+    for flag, description in platforms:
+        # Remove leading dashes for cleaner display
+        name = flag.lstrip("-")
+        help_lines.append(f"  - {name.upper():<10} {description}")
+
+    return "\n".join(help_lines)
+
+
+def get_hardware_args_group(parser):
+    """
+    Add hardware platform arguments to an argument parser.
+
+    Args:
+        parser: argparse.ArgumentParser instance
+
+    Returns:
+        The argument group for hardware platforms
+    """
+    hardware_group = parser.add_argument_group("Hardware Platform Options")
+
+    for flag, description in get_supported_hardware_platforms():
+        hardware_group.add_argument(flag, action="store_true", help=description)
+
+    return hardware_group
+
 
 def get_args():
-    """Parse command line arguments"""
-    parser = argparse.ArgumentParser(description="Test Operator")
+    """Parse command line arguments for operator testing"""
+    parser = argparse.ArgumentParser(
+        description="Test InfiniCore operators across multiple hardware platforms",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=f"""
+Examples:
+  # Run all tests on CPU only
+  python test_operator.py --cpu
+
+  # Run with benchmarking on NVIDIA GPU
+  python test_operator.py --nvidia --bench
+
+  # Run with debug mode on multiple devices
+  python test_operator.py --cpu --nvidia --debug
+
+  # Run performance profiling with custom iterations
+  python test_operator.py --nvidia --bench --num_prerun 50 --num_iterations 5000
+
+{get_hardware_help_text()}
+        """,
+    )
+
+    # Core testing options
     parser.add_argument(
         "--bench",
         action="store_true",
-        help="Whether to benchmark performance",
+        help="Enable performance benchmarking mode",
     )
     parser.add_argument(
         "--num_prerun",
         type=lambda x: max(0, int(x)),
         default=10,
-        help="Set the number of pre-runs before benchmarking. Default is 10.",
+        help="Number of warm-up runs before benchmarking (default: 10)",
     )
     parser.add_argument(
         "--num_iterations",
         type=lambda x: max(0, int(x)),
         default=1000,
-        help="Set the number of iterations for benchmarking. Default is 1000.",
+        help="Number of iterations for benchmarking (default: 1000)",
     )
     parser.add_argument(
         "--debug",
         action="store_true",
-        help="Whether to turn on debug mode.",
+        help="Enable debug mode for detailed tensor comparison",
     )
 
-    # Device options
-    device_group = parser.add_argument_group("Device options")
-    device_group.add_argument("--cpu", action="store_true", help="Run CPU test")
-    device_group.add_argument(
-        "--nvidia", action="store_true", help="Run NVIDIA GPU test"
-    )
-    device_group.add_argument(
-        "--cambricon", action="store_true", help="Run Cambricon MLU test"
-    )
-    device_group.add_argument(
-        "--ascend", action="store_true", help="Run ASCEND NPU test"
-    )
-    device_group.add_argument(
-        "--iluvatar", action="store_true", help="Run Iluvatar GPU test"
-    )
-    device_group.add_argument("--metax", action="store_true", help="Run METAX GPU test")
-    device_group.add_argument(
-        "--moore", action="store_true", help="Run MTHREADS GPU test"
-    )
-    device_group.add_argument(
-        "--kunlun", action="store_true", help="Run KUNLUN XPU test"
-    )
+    # Device options using shared hardware info
+    hardware_group = get_hardware_args_group(parser)
 
     return parser.parse_args()
 
 
 def get_test_devices(args):
     """
     Determine which devices to test based on command line arguments
+
+    Returns:
+        List[InfiniDeviceEnum]: List of devices to test
     """
     devices_to_test = []
 
+    # Check each hardware platform with proper dependency validation
     if args.cpu:
         devices_to_test.append(InfiniDeviceEnum.CPU)
+
     if args.nvidia:
-        devices_to_test.append(InfiniDeviceEnum.NVIDIA)
-    if args.iluvatar:
-        devices_to_test.append(InfiniDeviceEnum.ILUVATAR)
+        try:
+            import torch.cuda
+
+            devices_to_test.append(InfiniDeviceEnum.NVIDIA)
+        except ImportError:
+            print("Warning: CUDA not available, skipping NVIDIA tests")
+
     if args.cambricon:
         try:
             import torch_mlu
 
             devices_to_test.append(InfiniDeviceEnum.CAMBRICON)
         except ImportError:
             print("Warning: torch_mlu not available, skipping Cambricon tests")
+
     if args.ascend:
         try:
             import torch
@@ -82,10 +154,15 @@ def get_test_devices(args):
             devices_to_test.append(InfiniDeviceEnum.ASCEND)
         except ImportError:
             print("Warning: torch_npu not available, skipping Ascend tests")
+
     if args.metax:
-        import torch
+        try:
+            import torch
+
+            devices_to_test.append(InfiniDeviceEnum.METAX)
+        except ImportError:
+            print("Warning: Metax GPU support not available")
 
-        devices_to_test.append(InfiniDeviceEnum.METAX)
     if args.moore:
         try:
             import torch
@@ -94,6 +171,16 @@ def get_test_devices(args):
             devices_to_test.append(InfiniDeviceEnum.MOORE)
         except ImportError:
             print("Warning: torch_musa not available, skipping Moore tests")
+
+    if args.iluvatar:
+        try:
+            # Iluvatar GPU detection
+            import torch
+
+            devices_to_test.append(InfiniDeviceEnum.ILUVATAR)
+        except ImportError:
+            print("Warning: Iluvatar GPU support not available")
+
     if args.kunlun:
         try:
             import torch_xmlir
@@ -102,8 +189,17 @@ def get_test_devices(args):
         except ImportError:
             print("Warning: torch_xmlir not available, skipping Kunlun tests")
 
+    if args.hygon:
+        try:
+            import torch
+
+            devices_to_test.append(InfiniDeviceEnum.HYGON)
+        except ImportError:
+            print("Warning: Hygon DCU support not available")
+
     # Default to CPU if no devices specified
     if not devices_to_test:
         devices_to_test = [InfiniDeviceEnum.CPU]
+        print("No devices specified, defaulting to CPU")
 
     return devices_to_test

test/infinicore/framework/devices.py

@@ -22,7 +22,6 @@ class InfiniDeviceEnum:
     InfiniDeviceEnum.HYGON: "Hygon",
 }
 
-# Mapping that maps InfiniDeviceEnum to torch device string
 torch_device_map = {
     InfiniDeviceEnum.CPU: "cpu",
     InfiniDeviceEnum.NVIDIA: "cuda",