Skip to content

Commit 1bc9df4

Browse files
metaflowGoogle-ML-Automation
authored andcommitted
Updates LLVM usage to match [33fcd6acc755](llvm/llvm-project@33fcd6acc755) PiperOrigin-RevId: 698742870
1 parent f18df8f commit 1bc9df4

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

jax/experimental/mosaic/gpu/utils.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,12 @@ def globaltimer(kind: Literal["low", "high"] | None = None):
296296

297297

298298
def bytewidth(ty: ir.Type):
299+
# The actual width of TF32 is 19 bits. However, sinc we need to treat it as
300+
# 32 bits for compatibility reasons. TF32 used to be 32 bits wide in upstream
301+
# MLIR, but it changed in
302+
# https://github.com/llvm/llvm-project/commit/67a1fdb014790a38a205d28e1748634de34471dd.
303+
if ir.FloatTF32Type.isinstance(ty):
304+
return 4
299305
if ir.IntegerType.isinstance(ty):
300306
return ir.IntegerType(ty).width // 8
301307
if ir.FloatType.isinstance(ty):

0 commit comments

Comments
 (0)