github/ROCm
1
1
Fork 0
mirror of https://github.com/ROCm/ROCm.git synced 2026-04-05 03:01:17 -04:00
Code Issues Packages Projects Releases Wiki Activity

[FRONTEND][BACKEND] dd memory synchronization scope parameter to atomic ops. (#2562)

Browse Source 
Co-authored-by: Keren Zhou <kerenzhou@openai.com>
This commit is contained in:
Chris Jones
2023-10-31 02:18:27 +00:00
committed by GitHub
parent 70fca00b67
commit 2398b82f18
9 changed files with 128 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
include/triton/Dialect/Triton/IR/TritonAttrDefs.td
View File

@@ -66,6 +66,16 @@ def TT_AtomicRMWAttr : I32EnumAttr<
let cppNamespace = "::mlir::triton";
}
def TT_MemSyncScopeAttr : I32EnumAttr<
"MemSyncScope", "",
[
I32EnumAttrCase<"GPU", 1, "gpu">,
I32EnumAttrCase<"CTA", 2, "cta">,
I32EnumAttrCase<"SYSTEM", 3, "sys">,
]> {
let cppNamespace = "::mlir::triton";
}
// Program ID dimensions.
def TT_ProgramDim : I32EnumAttr<
"ProgramIDDim", "",

4
include/triton/Dialect/Triton/IR/TritonOps.td
View File

@@ -242,7 +242,7 @@ def TT_AtomicRMWOp : TT_Op<"atomic_rmw", [SameOperandsAndResultShape,
let arguments = (ins TT_AtomicRMWAttr:$atomic_rmw_op, TT_PtrLike:$ptr,
TT_Type:$val, Optional<TT_BoolLike>:$mask,
TT_MemSemanticAttr:$sem);
TT_MemSemanticAttr:$sem, TT_MemSyncScopeAttr:$scope);
let results = (outs TT_Type:$result);
}
@@ -264,7 +264,7 @@ def TT_AtomicCASOp : TT_Op<"atomic_cas", [MemoryEffects<[MemRead]>,
}];
let arguments = (ins TT_PtrLike:$ptr, TT_Type:$cmp, TT_Type:$val,
TT_MemSemanticAttr:$sem);
TT_MemSemanticAttr:$sem, TT_MemSyncScopeAttr:$scope);
let results = (outs TT_Type:$result);
}