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

[FRONTEND] Fixed up type cast in atomics codegen (#853)

Browse Source
This commit is contained in:
Philippe Tillet
2022-11-07 05:46:24 -08:00
committed by GitHub
parent 84ad215268
commit b6f15e214b
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/src/triton.cc
View File

@@ -1080,7 +1080,8 @@ void init_triton_ir(py::module &&m) {
[](mlir::OpBuilder &self, mlir::Value &ptr, mlir::Value &cmp, [](mlir::OpBuilder &self, mlir::Value &ptr, mlir::Value &cmp,
mlir::Value &val) -> mlir::Value { mlir::Value &val) -> mlir::Value {
auto loc = self.getUnknownLoc(); auto loc = self.getUnknownLoc();
auto ptrType = ptr.getType().dyn_cast<mlir::triton::PointerType>(); auto ptrType = mlir::getElementTypeOrSelf(ptr)
.cast<mlir::triton::PointerType>();
mlir::Type dstType = ptrType.getPointeeType(); mlir::Type dstType = ptrType.getPointeeType();
return self.create<mlir::triton::AtomicCASOp>(loc, dstType, ptr, return self.create<mlir::triton::AtomicCASOp>(loc, dstType, ptr,
cmp, val); cmp, val);
@@ -1090,7 +1091,8 @@ void init_triton_ir(py::module &&m) {
mlir::Value &ptr, mlir::Value &val, mlir::Value &ptr, mlir::Value &val,
mlir::Value &mask) -> mlir::Value { mlir::Value &mask) -> mlir::Value {
auto loc = self.getUnknownLoc(); auto loc = self.getUnknownLoc();
auto ptrType = ptr.getType().dyn_cast<mlir::triton::PointerType>(); auto ptrType = mlir::getElementTypeOrSelf(ptr)
.cast<mlir::triton::PointerType>();
mlir::Type dstType = ptrType.getPointeeType(); mlir::Type dstType = ptrType.getPointeeType();
return self.create<mlir::triton::AtomicRMWOp>(loc, dstType, rmwOp, return self.create<mlir::triton::AtomicRMWOp>(loc, dstType, rmwOp,
ptr, val, mask); ptr, val, mask);