#!/usr/bin/env python3 # coding: utf-8 # Copyright (c) 2025 Huawei Technologies Co., Ltd. # This program is free software, you can redistribute it and/or modify it under the terms and conditions of # CANN Open Software License Agreement Version 2.0 (the "License"). # Please refer to the License for details. You may not use this file except in compliance with the License. # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- """ Softmax Example for PyPTO This example demonstrates how to implement a softmax operation using PyPTO, including: - Manual softmax computation from basic operations - Dynamic axis marking for variable batch sizes - Tiling configuration for efficient execution - Loop-based processing for large tensors Softmax is a fundamental operation in neural networks, especially for attention mechanisms. """ import os import sys import argparse import pypto import torch import numpy as np from numpy.testing import assert_allclose def _peek_run_mode_from_argv(default: str = "npu") -> str: """Read run_mode early so module-level decorators can use it.""" for idx, arg in enumerate(sys.argv): if arg == "--run_mode" and idx + 1 < len(sys.argv): value = sys.argv[idx + 1] if value in ("npu", "sim"): return value if arg.startswith("--run_mode="): value = arg.split("=", 1)[1] if value in ("npu", "sim"): return value return default global_run_mode = pypto.RunMode.NPU if _peek_run_mode_from_argv("npu") == "sim": global_run_mode = pypto.RunMode.SIM def get_device_id(): """ Get and validate TILE_FWK_DEVICE_ID from environment variable. Returns: int: The device ID if valid, None otherwise. """ if 'TILE_FWK_DEVICE_ID' not in os.environ: print("Please set the environment variable TILE_FWK_DEVICE_ID before running:") print(" export TILE_FWK_DEVICE_ID=0") return None try: device_id = int(os.environ['TILE_FWK_DEVICE_ID']) return device_id except ValueError: print(f"ERROR: TILE_FWK_DEVICE_ID must be an integer, got: {os.environ['TILE_FWK_DEVICE_ID']}") return None def softmax_core(x: pypto.Tensor) -> pypto.Tensor: """ Core softmax computation: exp(x - max(x)) / sum(exp(x - max(x))). Parameters ---------- input_tensor : pypto.Tensor Input tensor to apply softmax to Returns ------- pypto.tensor Softmax normalized tensor """ row_max = pypto.amax(x, dim=-1, keepdim=True) sub = x - row_max exp = pypto.exp(sub) esum = pypto.sum(exp, dim=-1, keepdim=True) return exp / esum @pypto.frontend.jit(runtime_options={"run_mode": global_run_mode}) def softmax_kernel( input_tensor: pypto.Tensor([pypto.DYNAMIC, ...], pypto.DT_FP32), output_tensor: pypto.Tensor([pypto.DYNAMIC, ...], pypto.DT_FP32)): bs, seqlen, head, dim = input_tensor.shape tile_b = 1 # Process one batch at a time b_loop = bs // tile_b # Tiling shape setting for efficient execution pypto.set_vec_tile_shapes(1, 4, 1, 64) for idx in pypto.loop(0, b_loop, 1, name="LOOP_L0_bIdx", idx_name="idx"): b_offset = idx * tile_b b_offset_end = (idx + 1) * tile_b input_view = input_tensor[b_offset:b_offset_end, :seqlen, :head, :dim] softmax_out = softmax_core(input_view) output_tensor[b_offset:, ...] = softmax_out def test_softmax(device_id: int = None, dynamic: bool = True) -> None: device = f'npu:{device_id}' if global_run_mode == pypto.RunMode.NPU and device_id is not None else 'cpu' shape = (32, 32, 1, 256) x = torch.rand(shape, dtype=torch.float, device=device) y = torch.zeros(shape, dtype=torch.float, device=device) softmax_kernel(x, y) # default dim: -1 golden = torch.softmax(x, dim=-1).cpu() y = y.cpu() max_diff = np.abs(y.numpy() - golden.numpy()).max() print(f"Input shape: {x.shape}") print(f"Output shape: {y.shape}") print(f"Max difference: {max_diff:.6f}") if global_run_mode == pypto.RunMode.NPU: assert_allclose(np.array(y), np.array(golden), rtol=3e-3, atol=3e-3) print("✓ Softmax test passed") print() def main(): """Run softmax example. Usage: python softmax.py # Run example python softmax.py --list # List available examples """ parser = argparse.ArgumentParser( description="PyPTO Softmax Example", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: %(prog)s softmax::test_softmax Run the softmax::test_softmax example %(prog)s --list List all available examples """ ) parser.add_argument( 'example_id', type=str, nargs='?', help='Example ID to run (1). If not specified, the example will run.' ) parser.add_argument( '--list', action='store_true', help='List all available examples and exit' ) parser.add_argument( '--run_mode', type=str, nargs='?', default='npu', choices=["npu", "sim"], help='Run mode, supports npu and sim.' ) args = parser.parse_args() # Define available examples examples = { "softmax::test_softmax": { 'name': 'Softmax', 'description': 'Softmax implementation with dynamic batch size', 'function': test_softmax } } # List examples if requested if args.list: print("\n" + "=" * 60) print("Available Examples") print("=" * 60 + "\n") for ex_id, ex_info in sorted(examples.items()): print(f" ID: {ex_id}") print(f" name: {ex_info['name']}") print(f" description: {ex_info['description']}\n") return # Validate example ID if provided if args.example_id is not None: if args.example_id not in examples: print(f"ERROR: Invalid example ID: {args.example_id}") print(f"Valid example IDs are: {', '.join(map(str, sorted(examples.keys())))}") print("\nUse --list to see all available examples.") sys.exit(1) print("\n" + "=" * 60) print("PyPTO Softmax Example") print("=" * 60 + "\n") # Get and validate device ID (needed for NPU examples) device_id = None examples_to_run = [] if args.example_id is not None: # Run single example examples_to_run = [(args.example_id, examples[args.example_id])] else: # Run all examples examples_to_run = list(examples.items()) if args.run_mode == "npu": device_id = get_device_id() if device_id is None: return import torch_npu torch.npu.set_device(device_id) print("Running examples that require NPU hardware...") print("(Make sure CANN environment is configured and NPU is available)\n") try: for ex_id, ex_info in examples_to_run: print(f"Running Example {ex_id}: {ex_info['name']}") ex_info['function'](device_id) if len(examples_to_run) > 1: print("=" * 60) print("All softmax tests passed!") print("=" * 60) except Exception as e: print(f"\nError: {e}") raise if __name__ == "__main__": main()