[Test] Use knobs in test_core.py to set up env (#7902)

Affected: ``test_enable_fp_fusion``, ``test_override_arch``

See also: triton-lang/triton#7801

Signed-off-by: Ilya Veselov <[email protected]>

Author: joviliast

python/test/unit/language/test_core.py

@@ -6930,7 +6930,7 @@ def test_dot_max_num_imprecise_acc(M, N, K, BLOCK_M, BLOCK_N, BLOCK_K, in_type_s
 
 @pytest.mark.parametrize("enable_fp_fusion", [False, True])
 @pytest.mark.parametrize("default_override", [False, True])
-def test_enable_fp_fusion(enable_fp_fusion, default_override, device, monkeypatch):
+def test_enable_fp_fusion(enable_fp_fusion, default_override, device, fresh_knobs):
     # Sequential multiply add can be fused by backend
     @triton.jit
     def mul_add(data):
@@ -6939,7 +6939,7 @@ def mul_add(data):
 
     data = torch.randn((128, ), device=device, dtype=torch.float32)
     if default_override:
-        monkeypatch.setenv("TRITON_DEFAULT_FP_FUSION", "1" if enable_fp_fusion else "0")
+        fresh_knobs.language.default_fp_fusion = enable_fp_fusion
         h = mul_add.warmup(data, grid=(1, ))
     else:
         h = mul_add.warmup(data, grid=(1, ), enable_fp_fusion=enable_fp_fusion)
@@ -6957,7 +6957,7 @@ def mul_add(data):
 
 @pytest.mark.parametrize("arch", ["sm70", "sm80", "sm90", "gfx942", "gfx950", "gfx1200"])
 @pytest.mark.parametrize("env_var_override", [False, True])
-def test_override_arch(arch, env_var_override, device, monkeypatch):
+def test_override_arch(arch, env_var_override, device, fresh_knobs):
     if arch.startswith("sm") and not is_cuda():
         pytest.skip(f"{arch} arch only for CUDA")
     elif arch.startswith("gfx") and not is_hip():
@@ -6974,7 +6974,7 @@ def simple(data, out):
 
     if is_cuda():
         if env_var_override:
-            monkeypatch.setenv("TRITON_OVERRIDE_ARCH", arch)
+            fresh_knobs.runtime.override_arch = str(arch)
             h = simple.warmup(data, out, grid=(1, ))
         else:
             h = simple.warmup(data, out, arch=arch, grid=(1, ))
@@ -6984,7 +6984,7 @@ def simple(data, out):
         # For HIP, the generated kernel is a binary containing the final ISA. So we cannot run
         # them like CUDA side if the chip doesn't match. Here we just check generated ISA.
         if env_var_override:
-            monkeypatch.setenv("TRITON_OVERRIDE_ARCH", str(arch))
+            fresh_knobs.runtime.override_arch = str(arch)
             h = simple.warmup(data, out, grid=(1, ))
         else:
             h = simple.warmup(data, out, arch=arch, grid=(1, ))