[LoadStoreOpToLLVM] Broadcast the result for atomic rmw & cas ops (#5031)

dev-tomek · web-flow · commit d847b105b928 · 2025-09-04T12:04:29.000-04:00
This PR fixes imporper result broadcasting for atomic rmw & cas ops reported in #4879 16bit atomic cas test cases remain skipped due to #5025
diff --git a/python/test/unit/language/test_core.py b/python/test/unit/language/test_core.py
@@ -2013,8 +2013,6 @@ def kernel(I, O):
 @pytest.mark.parametrize("size", [1, 4, 16])
 @pytest.mark.parametrize("op", ["add", "cas"])
 def test_tensor_atomic_use_result(dtype_str, size, op, device):
-    if is_xpu():
-        pytest.skip("FIXME: issue #4879")
     if is_hip():
         pytest.skip(
             "HIP is broken because (1) it doesn't support thread predicate in atomic cas, and (2) it doesn't support"
diff --git a/scripts/skiplist/a770/language.txt b/scripts/skiplist/a770/language.txt
@@ -751,3 +751,7 @@ python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_batc
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_batched_gemm_3d_tma
 # test_tensor_atomic_add_access_patterns
 python/test/unit/language/test_core.py::test_tensor_atomic_add_access_patterns[shape128-random_no_duplication-3-1-float32]
+# https://github.com/intel/intel-xpu-backend-for-triton/issues/5025
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-1-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-4-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-16-float16]
diff --git a/scripts/skiplist/arl-h/language.txt b/scripts/skiplist/arl-h/language.txt
@@ -609,3 +609,7 @@ python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_redu
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-min]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-or]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-xor]
+# https://github.com/intel/intel-xpu-backend-for-triton/issues/5025
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-1-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-4-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-16-float16]
diff --git a/scripts/skiplist/arl-s/language.txt b/scripts/skiplist/arl-s/language.txt
@@ -609,3 +609,7 @@ python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_redu
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-min]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-or]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-xor]
+# https://github.com/intel/intel-xpu-backend-for-triton/issues/5025
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-1-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-4-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-16-float16]
diff --git a/scripts/skiplist/conda/language.txt b/scripts/skiplist/conda/language.txt
@@ -235,3 +235,7 @@ python/test/unit/language/test_core.py::test_const[if-False-False]
 python/test/unit/language/test_core.py::test_unroll_attr
 python/test/unit/language/test_decorator.py::test_triton_heuristic
 python/test/unit/language/test_core.py::test_constexpr_if_return
+# https://github.com/intel/intel-xpu-backend-for-triton/issues/5025
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-1-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-4-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-16-float16]
diff --git a/scripts/skiplist/default/language.txt b/scripts/skiplist/default/language.txt
@@ -86,3 +86,7 @@ python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_redu
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-min]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-or]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-xor]
+# https://github.com/intel/intel-xpu-backend-for-triton/issues/5025
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-1-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-4-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-16-float16]
diff --git a/scripts/skiplist/lts/language.txt b/scripts/skiplist/lts/language.txt
@@ -337,3 +337,7 @@ python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_redu
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-min]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-or]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-xor]
+# https://github.com/intel/intel-xpu-backend-for-triton/issues/5025
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-1-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-4-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-16-float16]
diff --git a/scripts/skiplist/mtl/language.txt b/scripts/skiplist/mtl/language.txt
@@ -386,3 +386,7 @@ python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_redu
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-min]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-or]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-xor]
+# https://github.com/intel/intel-xpu-backend-for-triton/issues/5025
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-1-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-4-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-16-float16]
diff --git a/scripts/skiplist/xe2/language.txt b/scripts/skiplist/xe2/language.txt
@@ -83,3 +83,7 @@ python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_redu
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-min]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-or]
 python/test/unit/language/test_tensor_descriptor.py::test_tensor_descriptor_reduce[8-32-host-1-uint32-xor]
+# https://github.com/intel/intel-xpu-backend-for-triton/issues/5025
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-1-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-4-float16]
+python/test/unit/language/test_core.py::test_tensor_atomic_use_result[cas-16-float16]
diff --git a/third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp b/third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp
@@ -3237,10 +3237,9 @@ struct AtomicCASOpConversion
     }
 
     if (tensorTy) {
-      Type structTy = getTypeConverter()->convertType(tensorTy);
-      Value resultStruct = packLLElements(loc, getTypeConverter(), resultVals,
-                                          rewriter, structTy);
-      rewriter.replaceOp(op, {resultStruct});
+      finalizeTensorAtomicResults(op, tensorTy, rewriter, resultVals,
+                                  valueElemTy, b, mask, targetInfo,
+                                  getTypeConverter());
     }
     return success();
   }
@@ -3407,10 +3406,9 @@ struct AtomicRMWOpConversion
     }
 
     if (tensorTy) {
-      Type structTy = getTypeConverter()->convertType(tensorTy);
-      Value resultStruct = packLLElements(loc, getTypeConverter(), resultVals,
-                                          rewriter, structTy);
-      rewriter.replaceOp(op, {resultStruct});
+      finalizeTensorAtomicResults(op, tensorTy, rewriter, resultVals,
+                                  valueElemTy, b, threadPred, targetInfo,
+                                  getTypeConverter());
     }
     return success();
   }

Original file line number	Diff line number	Diff line change
`@@ -3237,10 +3237,9 @@ struct AtomicCASOpConversion`
`3237`	`3237`	`}`
`3238`	`3238`
`3239`	`3239`	`if (tensorTy) {`
`3240`		`- Type structTy = getTypeConverter()->convertType(tensorTy);`
`3241`		`- Value resultStruct = packLLElements(loc, getTypeConverter(), resultVals,`
`3242`		`- rewriter, structTy);`
`3243`		`- rewriter.replaceOp(op, {resultStruct});`
	`3240`	`+ finalizeTensorAtomicResults(op, tensorTy, rewriter, resultVals,`
	`3241`	`+ valueElemTy, b, mask, targetInfo,`
	`3242`	`+ getTypeConverter());`
`3244`	`3243`	`}`
`3245`	`3244`	`return success();`
`3246`	`3245`	`}`
`@@ -3407,10 +3406,9 @@ struct AtomicRMWOpConversion`
`3407`	`3406`	`}`
`3408`	`3407`
`3409`	`3408`	`if (tensorTy) {`
`3410`		`- Type structTy = getTypeConverter()->convertType(tensorTy);`
`3411`		`- Value resultStruct = packLLElements(loc, getTypeConverter(), resultVals,`
`3412`		`- rewriter, structTy);`
`3413`		`- rewriter.replaceOp(op, {resultStruct});`
	`3409`	`+ finalizeTensorAtomicResults(op, tensorTy, rewriter, resultVals,`
	`3410`	`+ valueElemTy, b, threadPred, targetInfo,`
	`3411`	`+ getTypeConverter());`
`3414`	`3412`	`}`
`3415`	`3413`	`return success();`
`3416`	`3414`	`}`