Skip to content

[Draft] Support the globaltimer and smid on Intel Arch #4816

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[Draft] Support the globaltimer and smid on Intel Arch #4816

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 14 additions & 6 deletions python/test/unit/language/test_core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6169,7 +6169,6 @@ def kernel(Out):


def test_globaltimer(device):
check_cuda_or_hip(device)
if is_hip():
pytest.skip("test_globaltimer is flaky on AMD GPUs")

Expand All @@ -6195,6 +6194,8 @@ def kernel(Out1, Out2, func: tl.constexpr):
assert out2[1] - out2[0] > 0
if is_cuda():
assert h.asm["ptx"].count("%globaltimer") == 2
elif is_xpu():
assert h.asm["llir"].count("%tsc") == 2
else:
target_arch = triton.runtime.driver.active.get_current_target().arch
if "gfx11" in target_arch or "gfx12" in target_arch:
Expand All @@ -6206,16 +6207,23 @@ def kernel(Out1, Out2, func: tl.constexpr):
def test_smid(device):
if is_hip():
pytest.skip("test_smid is not supported in HIP")
check_cuda_or_hip(device)

@triton.jit
def kernel(Out):
tl.store(Out + tl.program_id(0), tl.extra.intel.smid())
def kernel(Out, func: tl.constexpr):
tl.store(Out + tl.program_id(0), func())

if is_cuda():
func = tl.extra.cuda.smid
elif is_xpu():
func = tl.extra.intel.smid

out = to_triton(np.zeros((1024, ), dtype=np.int32), device=device)
h = kernel[(out.shape[0], )](out)
h = kernel[(out.shape[0], )](out, func)
assert out.sort()[0].unique().shape[0] > 0
assert h.asm["ptx"].count("%smid") == 1
if is_cuda():
assert h.asm["ptx"].count("%smid") == 1
elif is_xpu():
assert h.asm["llir"].count("%sr0") == 1


# -----------------------
Expand Down
14 changes: 10 additions & 4 deletions third_party/intel/language/intel/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,20 @@

@core.extern
def globaltimer(_semantic=None):
return core.inline_asm_elementwise("mov.u64 $0, %globaltimer;", "=l", [], dtype=core.int64, is_pure=False, pack=1,
_semantic=_semantic)
return core.inline_asm_elementwise(
"""{\n .decl globaltimer v_type=G type=ud num_elts=2 align=qword alias=<$0, 0> \n"""
""" mov (M1_NM, 2) globaltimer(0, 0)<1> %tsc(0,0)<1;1,0> \n}""", "=rw.u", [], dtype=core.uint64, is_pure=False,
pack=1, _semantic=_semantic)
Comment on lines +6 to +9
Copy link
Preview

Copilot AI Jul 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nitpick] The inline assembly string is complex and split across multiple lines with embedded newlines and escape sequences. Consider extracting this assembly code to a constant or using a more readable multiline string format to improve maintainability.

Suggested change
return core.inline_asm_elementwise(
"""{\n .decl globaltimer v_type=G type=ud num_elts=2 align=qword alias=<$0, 0> \n"""
""" mov (M1_NM, 2) globaltimer(0, 0)<1> %tsc(0,0)<1;1,0> \n}""", "=rw.u", [], dtype=core.uint64, is_pure=False,
pack=1, _semantic=_semantic)
GLOBALTIMER_ASM = (
"{\n"
" .decl globaltimer v_type=G type=ud num_elts=2 align=qword alias=<$0, 0> \n"
" mov (M1_NM, 2) globaltimer(0, 0)<1> %tsc(0,0)<1;1,0> \n"
"}"
)
return core.inline_asm_elementwise(
GLOBALTIMER_ASM, "=rw.u", [], dtype=core.uint64, is_pure=False, pack=1, _semantic=_semantic
)

Copilot uses AI. Check for mistakes.



@core.extern
def smid(_semantic=None):
return core.inline_asm_elementwise("mov.u32 $0, %smid;", "=r", [], dtype=core.int32, is_pure=True, pack=1,
_semantic=_semantic)
sr = core.inline_asm_elementwise("mov (M1_NM, 1) $0(0, 0)<1> %sr0(0,0)<0;1,0>", "=rw.u", [], dtype=core.uint32,
is_pure=True, pack=1, _semantic=_semantic)
pos: core.constexpr = core.constexpr(9)
Copy link
Preview

Copilot AI Jul 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The magic number 9 should be documented or extracted to a named constant to explain what bit position it represents in the status register.

Suggested change
pos: core.constexpr = core.constexpr(9)
pos: core.constexpr = core.constexpr(STATUS_REGISTER_BIT_POSITION)

Copilot uses AI. Check for mistakes.

subslice_mask: core.constexpr = core.constexpr((1 << 11) - 1)
Copy link
Preview

Copilot AI Jul 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The magic number 11 should be documented or extracted to a named constant to explain the bit width of the subslice mask.

Suggested change
subslice_mask: core.constexpr = core.constexpr((1 << 11) - 1)
SUBSLICE_MASK_BIT_WIDTH = 11 # Bit width of the subslice mask
subslice_mask: core.constexpr = core.constexpr((1 << SUBSLICE_MASK_BIT_WIDTH) - 1)

Copilot uses AI. Check for mistakes.

return sr.__and__(subslice_mask, _semantic=_semantic).__rshift__(pos, _semantic=_semantic)



@core.builtin
Expand Down
Loading