mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 20:36:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

Remove MemoryEffects annotations from async_{load/store} ops

Browse Source 
The annotation on async_load didn't indicate its write to SMEM, allowing it
to be DCEd by MLIR canonicalization. We don't get much mileage out of those
annotations, so let's just delete them for simplicity.

PiperOrigin-RevId: 731003033
This commit is contained in:
Adam Paszke 2025-02-25 13:14:17 -08:00 committed by jax authors
parent 03e2c888e2
commit cb7402f6de
2 changed files with 2 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
jaxlib/mosaic/dialect/gpu/mosaic_gpu.h
View File

@ -46,10 +46,6 @@ namespace mosaic_gpu {
using Memref = ::mlir::TypedValue<::mlir::MemRefType>;
using Pointer = ::mlir::TypedValue<::mlir::LLVM::LLVMPointerType>;
struct GlobalMemory : public mlir::SideEffects::Resource::Base<GlobalMemory> {
llvm::StringRef getName() final { return "<GlobalMemory>"; }
};
constexpr absl::string_view kRuntimeTmaDescriptorInitializerName =
"mosaic_gpu_init_tma_desc";

6
jaxlib/mosaic/dialect/gpu/mosaic_gpu.td
View File

@ -243,10 +243,8 @@ def LayoutAttr : MosaicGPU_Attr<"Layout", "layout",
let assemblyFormat = "`<` $num_dimensions `,` $transforms `>`";
}
def GlobalMemory : Resource<"::mosaic_gpu::GlobalMemory">;
def MosaicGPU_AsyncLoadOp : Op<MosaicGPU_Dialect, "async_load",
[AttrSizedOperandSegments, MemoryEffects<[MemRead<GlobalMemory>]>]> {
[AttrSizedOperandSegments]> {
let summary = "Schedules an async load of a MemRef from GMEM to SMEM";
let description = [{
Schedules an async copy of the contents of the `source` MemRef in GMEM to
@ -313,7 +311,7 @@ def MosaicGPU_AsyncLoadOp : Op<MosaicGPU_Dialect, "async_load",
}
def MosaicGPU_AsyncStoreOp : Op<MosaicGPU_Dialect, "async_store",
[AttrSizedOperandSegments, MemoryEffects<[MemWrite<GlobalMemory>]>]> {
[AttrSizedOperandSegments]> {
let summary = "Schedules an async store of a MemRef from SMEM to GMEM";
let description = [{
Schedules an async store of the contents of the `source` MemRef in SMEM to