Low-Level XPU Local Atomic Enhancement for Add & CAS

[yucai-intel]

This PR focuses on the src/ATen/native/xpu/sycl/Atomics.h file, aiming to fully implement and enhance atomic operations on Shared Local Memory SMEM to support performance optimizations in upper-layer kernels.

1. Implementation of Local AtomicCAS (Compare-and-Swap) Mechanism

• The PR introduces generic AtomicCASInteger and AtomicCASFP structs.
• It implements CAS operations on SMEM for all major data types, including int64, float, and Half.
• The implementation uses either native SYCL CAS Soft RMW loops to ensure correctness across different data widths.

2. Completeness and Correction of Local AtomicAdd Operations

• The PR provides direct SYCL implementations for atomicAddLocal for basic types like float, double, and int.
• It ensures the correct accumulation of half-precision floats like Half and BFloat16 on SMEM using a CAS loop mechanism.

This foundational work is crucial for enabling high-performance caching logic in operators like index_add.

[Copilot]

Pull request overview

This PR enhances atomic operations on XPU Shared Local Memory (SMEM) by implementing Compare-and-Swap (CAS) operations and completing local atomic add support for various data types. The changes enable high-performance caching logic in operators like index_add by providing foundation-level atomic primitives.

Key Changes:

• Introduced generic AtomicCASInteger and AtomicCASFP template structures supporting CAS operations on local memory for integer and floating-point types (including Half/BFloat16)
• Completed atomicAddLocal implementations for basic types (float, double, int variants) and half-precision types using CAS-based loops
• Fixed macro naming inconsistency by renaming SYCL_ATOMIC_INTEGER_LOCAL output from atomic##NAME to atomic##NAME##Local

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

This line calls fetch_add in atomicMax function, but should call fetch_max instead. The incorrect operation will not produce maximum values as intended.

[Copilot]

C-style pointer casts for type punning are undefined behavior in C++. Use std::bit_cast (C++20) or memcpy for safe type reinterpretation between float and unsigned int.

[Copilot]

C-style pointer casts for type punning are undefined behavior in C++. Use std::bit_cast (C++20) or memcpy for safe type reinterpretation between double and unsigned long long.

[Copilot]

C-style pointer cast for type punning is undefined behavior. Use std::bit_cast (C++20) or memcpy to safely reinterpret unsigned int bits as type T (float).

[Copilot]

C-style pointer cast for type punning is undefined behavior. Use std::bit_cast (C++20) or memcpy to safely reinterpret unsigned long long bits as type T (double).

[CuiYifeng]

Can you explain the motivation behind declaring at::Half hsum in the loop body?

[CuiYifeng]

I also wonder if such casts will cause undefined behavior.