gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[TRITON-MLIR][BACKEND]fix atomic_rmw for vector (#966)

Browse Source 
Co-authored-by: dongdongl <dongdongl@nvidia.com>
This commit is contained in:
Dongdong Li
2022-12-08 20:03:40 +08:00
committed by GitHub
parent f0885e9caf
commit c7cf9c6a32
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
lib/Conversion/TritonGPUToLLVM/TritonGPUToLLVM.cpp
View File

@@ -4420,6 +4420,7 @@ struct AtomicRMWOpConversion
rewriter, loc, getTypeConverter()->getIndexType(), ii); rewriter, loc, getTypeConverter()->getIndexType(), ii);
rmwVal = insert_element(vecTy, rmwVal, valElements[i + ii], iiVal); rmwVal = insert_element(vecTy, rmwVal, valElements[i + ii], iiVal);
} }
Value rmwPtr = ptrElements[i]; Value rmwPtr = ptrElements[i];
Value rmwMask = maskElements[i]; Value rmwMask = maskElements[i];
rmwMask = and_(rmwMask, mask); rmwMask = and_(rmwMask, mask);
@@ -4477,9 +4478,10 @@ struct AtomicRMWOpConversion
atom.o(rmwOp).o(sTy); atom.o(rmwOp).o(sTy);
if (valueTy) { if (valueTy) {
atom(dstOpr, ptrOpr, valOpr).predicate(rmwMask); atom(dstOpr, ptrOpr, valOpr).predicate(rmwMask);
auto ret = ptxBuilderAtomicRMW.launch(rewriter, loc, valueElemTy); auto retType = vec == 1 ? valueElemTy : vecTy;
auto ret = ptxBuilderAtomicRMW.launch(rewriter, loc, retType);
for (int ii = 0; ii < vec; ++ii) { for (int ii = 0; ii < vec; ++ii) {
resultVals[i * vec + ii] = resultVals[i + ii] =
vec == 1 ? ret : extract_element(valueElemTy, ret, idx_val(ii)); vec == 1 ? ret : extract_element(valueElemTy, ret, idx_val(ii));
} }
} else { } else {