Use inline VISA to optimize horizontal batched subgroup reduce

[chengjunlu]

Use inline VISA to optimize horizontal batched subgroup reduce.

1. Support the SIMD reduce on ATS for threads_per_warp=16
2. Support the SIMD reduce on PVC for threads_per_warp=16 or 32.
3. There is limitation that:

• Only supports the float32 and float16 for now.
• Only supports the sum and max for now.

Run the unit test CI.

[Copilot]

Pull Request Overview

This PR introduces an experimental inline VISA mechanism to optimize horizontal batched subgroup reduce in the Intel backend while providing stub implementations for NVIDIA and AMD backends. Key changes include:

• Adding a new warpBatchReduce function implementation with inline VISA in Intel’s TargetInfo.cpp.
• Updating header files across Intel, NVIDIA, and AMD backends and the base interface to declare/open up the new function.
• Integrating the new warpBatchReduce call into ReduceOpToLLVM.cpp for early returns when applicable.

Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TargetInfo.h	Add stub implementation for warpBatchReduce returning false.
third_party/intel/lib/TritonIntelGPUToLLVM/TargetInfo.h	Declare the new warpBatchReduce function.
third_party/intel/lib/TritonIntelGPUToLLVM/TargetInfo.cpp	Implement experimental inline VISA-based warpBatchReduce logic.
third_party/intel/lib/TritonIntelGPUToLLVM/ReduceOpToLLVM.cpp	Integrate warpBatchReduce into reduce op conversion.
third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.h	Add stub implementation for warpBatchReduce returning false.
include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h	Add pure virtual declaration for warpBatchReduce.

[Copilot]

[nitpick] Iterating over 'acc' using 'auto it' and then accessing 'acc[key]' results in redundant lookups; consider using structured bindings (e.g., 'for (auto &pair : acc)') to improve clarity and efficiency.

Suggested change

	for (auto it : acc) {
	const SmallVector<unsigned> &key = it.first;
	SmallVector<Value> &val = acc[key];
	for (auto &[key, val] : acc) {

[Copilot]

[nitpick] The hard-coded check for warpSize == 16 limits the function to experimental scenarios; consider adding a comment or an assert to clarify the dependency on this constraint.

Suggested change

	// It is only experimental code supports threads_per_warp=16
	// This code is experimental and currently supports only threads_per_warp=16.
	assert(warpSize == 16 && "This experimental code supports only warpSize of 16.");

[alexbaden]

Where is the GitHub issue for this work?

[alexbaden]

Do we need to add the method to the global target info if we are the only ones using it, inside files we control?

[chengjunlu]

It is still under investigating how to implement this to upstream.

Maybe we can use the in-tree MLIR ops: https://mlir.llvm.org/docs/Dialects/GPU/#gpusubgroup_reduce-gpusubgroupreduceop

[chengjunlu]

Where is the GitHub issue for this work?

#3310

[chengjunlu]

This is a large PR which is going to be split into several small ones.

1. Add SIMD reduce utils in XeAsmFormat.cpp and unit test.
2. Add a new placeholder in the ReduceOpLowering.cpp to override the default reduce with-in warp for multiple inputs.
3. Enable the SIMD reduce by default.

[chengjunlu]

The functionality of simd reduce is ready and tested.
The performance improvements is not expected on attention kernel for now. Need to further investigate.

[whitneywhtsang]

Due to #4171 (comment), moving this PR to draft.