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

[Triton-MLIR][FRONTEND] [BACKEND] fix atomics (#879)

Browse Source 
minor fix to backend and frontend of atomics, we can pass 1 test without
mask and the shape aligned with CTA size now

Co-authored-by: dongdongl <dongdongl@nvidia.com>
This commit is contained in:
donproc
2022-11-16 12:25:15 +08:00
committed by GitHub
parent 37f5846280
commit 5eee738df7
4 changed files with 31 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/src/triton.cc
View File

@@ -1086,9 +1086,7 @@ void init_triton_ir(py::module &&m) {
mlir::Value &ptr, mlir::Value &val,
mlir::Value &mask) -> mlir::Value {
auto loc = self.getUnknownLoc();
auto ptrType = mlir::getElementTypeOrSelf(ptr)
.cast<mlir::triton::PointerType>();
mlir::Type dstType = ptrType.getPointeeType();
mlir::Type dstType = val.getType();
return self.create<mlir::triton::AtomicRMWOp>(loc, dstType, rmwOp,
ptr, val, mask);
})