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

[Mosaic] Allow users to opt out of window prefetching

Browse Source 
PiperOrigin-RevId: 552797922
This commit is contained in:
Adam Paszke 2023-08-01 07:36:23 -07:00 committed by jax authors
parent b8019dcc9e
commit 716f4f8119
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
jaxlib/mosaic/dialect/tpu/tpu.td
View File

@ -111,8 +111,9 @@ def TPU_TiledLayoutAttr
}
def TPU_MemorySpace : I32EnumAttr<"MemorySpace", "Memory space", [
I32EnumAttrCase<"vmem", 0>,
I32EnumAttrCase<"smem", 1>
I32EnumAttrCase<"kAny", 4294967295, "any">,
I32EnumAttrCase<"vmem", 0, "vmem">,
I32EnumAttrCase<"smem", 1, "smem">
]> {
let genSpecializedAttr = 0;
let cppNamespace = "::mlir::tpu";