Background

When lowering masked store operations in the compiler, masked store is currently expanded into a load → select → store read-modify-write (RMW) sequence.

This transformation is non-atomic and introduces data races when multiple parallel tasks write to overlapping regions of the same output tensor, leading to non-deterministic results and silent precision errors.

Problem Example & Flowchart

Current Lowering (Buggy RMW Flow)

# Original kernel code (discrete mask store)
tl.store(ptr, val, mask=discrete_mask)

       ↓ Compiler Lowering

# Compiler lowers to unsafe RMW sequence
old_val = tl.load(ptr)          # Unprotected parallel read
selected = select(discrete_mask, val, old_val)
tl.store(ptr, selected)         # Unprotected parallel write

Data Race Flowchart (Parallel Execution)

Task A (Thread 0)       Task B (Thread 1)
    │                       │
    ├─ load (old = 0)       │
    │                       ├─ load (old = 0)
    │                       ├─ select → val = 5
    │                       ├─ store → mem = 5
    ├─ select → val = 0     │
    └─ store → mem = 0      │
            ↓
Final value: 0 (incorrect, non-deterministic)

• Both threads load the same initial value before either write completes
• The final result depends entirely on thread scheduling order
• Overlapped regions produce unstable/incorrect outputs → precision errors

Root Cause

1. Discrete masked store is lowered to a non-atomic three-step RMW
2. No synchronization mechanism for overlapping parallel writes
3. The intermediate load step reads stale/conflicting values
4. Write operations are not mutually exclusive, causing data corruption

Solution Requirements

We need two safe replacement methods:

1. Synchronized RMW (Default Fix)

Wrap the RMW sequence with block-level synchronization locks to guarantee atomicity for overlapping write scenarios.
Controlled by a compile flag enable_sync_block_lock (default: False):

• Default enabled (True) ensures data precision and atomicity for all overlapping write scenarios, eliminating race conditions.
• Can be explicitly set to False to disable synchronization for performance optimization in scenarios where mask overlap is guaranteed not to occur.

// Safe synchronized RMW flow
SyncBlockLock();
  old_val = load(ptr);
  selected = select(discrete_mask, val, old_val);
  store(ptr, selected);
SyncBlockUnlock();

Data Race Flowchart (Serial Equivalent)

Initial value: 0

Task A runs first
    │
    ├─ load → old = 0
    │
    ├─ select → 0
    └─ store → mem = 0

Task A runs after
    │
    ├─ load → old = 0
    │
    ├─ select → 5
    └─ store → mem = 5

Final value: 5 (CORRECT)

Advantages:

• Ensures mutual exclusion within block scope
• Eliminates data races for overlapping writes
• Guarantees deterministic and correct results
• Preserves full compatibility with existing masked store semantics

2. Direct Indirect Store (Optimal Fix)

For applicable scenarios, completely eliminate the RMW sequence and lower directly to hardware-supported indirect store operations when force_simt_template=True:

// Safe: no load step, zero data race risk
indirect_store(ptr, val, discrete_mask);

Advantages:

• Removes the unsafe load operation entirely
• Native hardware discrete mask support
• Optimal performance and full correctness
• Eliminates data race conditions by design

Architectural Decision

Decision: Introduce enable_sync_block_lock (default: False)

Rationale

• Current Triton-Ascend mask analysis cannot automatically detect overlapping mask scenarios in masked store operations.
• In practice, most use cases of masked store do not involve overlapping mask regions; setting the default value to False ensures optimal performance for the majority of scenarios.
• Users can explicitly set enable_sync_block_lock=True to enable block-level synchronization when they cannot guarantee non-overlapping masks, thereby ensuring data correctness and eliminating race conditions.

Future Plan
Once the mask analysis module is enhanced to automatically identify non-overlapping mask regions in store operations, this flag will become redundant and can be removed entirely.