Skip to content

Commit

Permalink
[TEST] Remove unnecessary capability query in test_subproc.py (#4432)
Browse files Browse the repository at this point in the history
It's an unused parameter. Seems like it was used in the first iteration
of this test, but that changed along the way. Removing this query makes
the test portable across backends.
  • Loading branch information
int3 committed Jul 30, 2024
1 parent 86f0822 commit 72734f0
Showing 1 changed file with 4 additions and 12 deletions.
16 changes: 4 additions & 12 deletions python/test/unit/runtime/test_subproc.py
Original file line number Diff line number Diff line change
@@ -1,16 +1,14 @@
import multiprocessing
import shutil

import torch

import triton
import triton.language as tl
from triton.compiler import ASTSource

target = triton.runtime.driver.active.get_current_target()


def compile_fn(attrs, capability):
def compile_fn(attrs):

@triton.jit
def kernel_sub(a, b, o, N: tl.constexpr):
Expand All @@ -27,18 +25,15 @@ def kernel_sub(a, b, o, N: tl.constexpr):


def test_compile_in_subproc() -> None:
major, minor = torch.cuda.get_device_capability(0)
cc = major * 10 + minor
config = triton.compiler.AttrsDescriptor(tuple(range(4)), ())

multiprocessing.set_start_method('fork')
proc = multiprocessing.Process(target=compile_fn, args=(config, cc))
proc = multiprocessing.Process(target=compile_fn, args=(config, ))
proc.start()
proc.join()
assert proc.exitcode == 0


def compile_fn_dot(attrs, capability):
def compile_fn_dot(attrs):

@triton.jit
def kernel_dot(Z):
Expand All @@ -52,12 +47,9 @@ def kernel_dot(Z):


def test_compile_in_forked_subproc(fresh_triton_cache) -> None:
major, minor = torch.cuda.get_device_capability(0)
capability = major * 10 + minor
config = triton.compiler.AttrsDescriptor(tuple(range(1)), ())

assert multiprocessing.get_start_method() == 'fork'
proc = multiprocessing.Process(target=compile_fn_dot, args=(config, capability))
proc = multiprocessing.Process(target=compile_fn_dot, args=(config, ))
proc.start()
proc.join()
assert proc.exitcode == 0
Expand Down

0 comments on commit 72734f0

Please sign in to comment.