revise test

Author: rosenrodt

tests/unittest/_torch/misc/test_autotuner.py

@@ -215,34 +215,36 @@ def bucket_mapper(x: int, inflate_factor: int) -> int:
         "test_runtime_bucket_mapping")
 
     # Extract the first dimension of the first input shape from each cache key
-    cached_bucket_sizes = sorted([key[3][0][0] for key in cache_entries.keys()])
-    expected_buckets = [1, 2, 4, 8, 16, 32]
-    assert cached_bucket_sizes == expected_buckets, \
-        f"Expect {expected_buckets}, got {cached_bucket_sizes}"
+    assert len(cache_entries) == 6, \
+        f"Expected 6 cache entries (buckets 1, 2, 4, 8, 16, 32), got {len(cache_entries)}"
 
     # Test runtime mapping: input size should be mapped via map_to_runtime_buckets
     # to find the correct tuning bucket
     test_cases = [
-        # input_size, expected_mapped_bucket via round_rule
-        (4, 1),  # round_rule(4) = 4 // 4 = 1
-        (8, 2),  # round_rule(8) = 8 // 4 = 2
-        (16, 4),  # round_rule(16) = 16 // 4 = 4
-        (32, 8),  # round_rule(32) = 32 // 4 = 8
-        (64, 16),  # round_rule(64) = 64 // 4 = 16
-        (128, 32),  # round_rule(128) = 128 // 4 = 32
+        # size 4 maps to bucket 4//4 = 1, tactic 0 (1 <= M // 2)
+        (4, 1, 0),
+        # size 8 maps to bucket 8//4 = 2, tactic 0 (2 <= M // 2)
+        (8, 2, 0),
+        # size 16 maps to bucket 16//4 = 4, tactic 0 (4 <= M // 2)
+        (16, 4, 0),
+        # size 32 maps to bucket 32//4 = 8, tactic 0 (8 <= M // 2)
+        (32, 8, 0),
+        # size 64 maps to bucket 64//4 = 16, tactic 0 (16 <= M // 2)
+        (64, 16, 0),
+        # size 128 maps to bucket 128//4 = 32, tactic 1 (32 > M // 2)
+        (128, 32, 1),
+        # size 256 maps to bucket 256//4 = 64, tactic -1 (64 > M)
+        (256, 64, -1),
     ]
 
-    for input_size, expected_bucket in test_cases:
-        # Verify the round_rule mapping
-        assert bucket_mapper(input_size, inflate_factor=4) == expected_bucket, \
-            f"bucket_mapper({input_size}, inflate_factor=4) should be {expected_bucket}, got {bucket_mapper(input_size, inflate_factor=4)}"
-
+    for input_size, expected_bucket, expected_tactic in test_cases:
         # Verify cache lookup succeeds with the mapped bucket
         x = torch.randn(input_size, 64)
         runner, tactic = tuner.choose_one("test_runtime_bucket_mapping",
                                           runners, tuning_config, [x, w])
-        assert tactic != -1, \
-            f"Cache miss for input_size={input_size}, expected to map to bucket {expected_bucket}"
+        assert (
+            tactic == expected_tactic
+        ), f"Cache mismatch for input_size={input_size}, expected to map to bucket {expected_tactic} but got {tactic}"
 
 
 def test_autotuner_try_block():