已合并
classify elementwise symbolic workloads #42677
luqichao创建于 7月25日
classify elementwise symbolic workloads #42677
已合并
共 4 个文件变更+284-3
| @@ -21,6 +21,7 @@ import unittest | |||
| 21 | import torch | 21 | import torch |
| 22 | import torch.nn.functional as F | 22 | import torch.nn.functional as F |
| 23 | from torch._dynamo.testing import CompileCounterWithBackend | 23 | from torch._dynamo.testing import CompileCounterWithBackend |
| 24 | +from torch._inductor.utils import run_and_get_code | ||
| 24 | from torch.testing._internal.common_utils import ( | 25 | from torch.testing._internal.common_utils import ( |
| 25 | instantiate_parametrized_tests, | 26 | instantiate_parametrized_tests, |
| 26 | parametrize, | 27 | parametrize, |
| @@ -528,6 +529,106 @@ class TestFusionDynamicShapes(DynamicShapeTestMixin, TestCase): | |||
| 528 | self._assert_compile_count() | 529 | self._assert_compile_count() |
| 529 | 530 | ||
| 530 | 531 | ||
| 532 | +# ============================================================ | ||
| 533 | +# Symbolic grouped-autotune workload classification | ||
| 534 | +# ============================================================ | ||
| 535 | + | ||
| 536 | +class TestSymbolicGroupElementwise(TestCase): | ||
| 537 | + def setUp(self): | ||
| 538 | + super().setUp() | ||
| 539 | + import torch_npu._inductor.config as npu_config | ||
| 540 | + | ||
| 541 | + self.npu_config = npu_config | ||
| 542 | + self.prev_group_autotune = npu_config.enable_symbolic_shape_group_autotune | ||
| 543 | + npu_config.enable_symbolic_shape_group_autotune = True | ||
| 544 | + torch._dynamo.reset() | ||
| 545 | + | ||
| 546 | + def tearDown(self): | ||
| 547 | + self.npu_config.enable_symbolic_shape_group_autotune = self.prev_group_autotune | ||
| 548 | + torch._dynamo.reset() | ||
| 549 | + super().tearDown() | ||
| 550 | + | ||
| 551 | + def _run_and_check(self, fn, inputs, next_inputs, expected_workload): | ||
| 552 | + for current, next_value in zip(inputs, next_inputs): | ||
| 553 | + if isinstance(current, torch.Tensor) and current.ndim and current.shape[0] != next_value.shape[0]: | ||
| 554 | + torch._dynamo.mark_dynamic(current, 0) | ||
| 555 | + | ||
| 556 | + expected = fn(*inputs) | ||
| 557 | + compiled = torch.compile(fn, backend="inductor") | ||
| 558 | + actual, codes = run_and_get_code(compiled, *inputs) | ||
| 559 | + torch.testing.assert_close(actual, expected) | ||
| 560 | + | ||
| 561 | + next_expected = fn(*next_inputs) | ||
| 562 | + next_actual = compiled(*next_inputs) | ||
| 563 | + torch.testing.assert_close(next_actual, next_expected) | ||
| 564 | + | ||
| 565 | + expected_metadata = f"'group_workload': {expected_workload!r}" | ||
| 566 | + matching_codes = [ | ||
| 567 | + code | ||
| 568 | + for code in codes | ||
| 569 | + if "'group_enabled': True" in code | ||
| 570 | + and "'group_template': 'pointwise'" in code | ||
| 571 | + and expected_metadata in code | ||
| 572 | + ] | ||
| 573 | + self.assertTrue( | ||
| 574 | + matching_codes, | ||
| 575 | + f"Expected pointwise group metadata with {expected_metadata}, got:\n{codes}", | ||
| 576 | + ) | ||
| 577 | + | ||
| 578 | + def test_basic_elementwise_workload(self): | ||
| 579 | + def fn(x, y): | ||
| 580 | + return torch.relu(x + y) * 0.5 | ||
| 581 | + | ||
| 582 | + inputs = ( | ||
| 583 | + torch.randn((257, 1031), device=device), | ||
| 584 | + torch.randn((257, 1031), device=device), | ||
| 585 | + ) | ||
| 586 | + next_inputs = ( | ||
| 587 | + torch.randn((263, 1031), device=device), | ||
| 588 | + torch.randn((263, 1031), device=device), | ||
| 589 | + ) | ||
| 590 | + self._run_and_check(fn, inputs, next_inputs, "elementwise") | ||
| 591 | + | ||
| 592 | + def test_broadcast_is_not_elementwise_workload(self): | ||
| 593 | + def fn(x, bias): | ||
| 594 | + return torch.relu(x + bias) * 0.5 | ||
| 595 | + | ||
| 596 | + bias = torch.randn((1031,), device=device) | ||
| 597 | + inputs = (torch.randn((257, 1031), device=device), bias) | ||
| 598 | + next_inputs = (torch.randn((263, 1031), device=device), bias) | ||
| 599 | + self._run_and_check(fn, inputs, next_inputs, None) | ||
| 600 | + | ||
| 601 | + def test_consumed_full_is_elementwise_workload(self): | ||
| 602 | + def fn(x): | ||
| 603 | + generated = torch.full_like(x, 2.0) | ||
| 604 | + return torch.relu(x + generated) | ||
| 605 | + | ||
| 606 | + inputs = (torch.randn((257, 1031), device=device),) | ||
| 607 | + next_inputs = (torch.randn((263, 1031), device=device),) | ||
| 608 | + self._run_and_check(fn, inputs, next_inputs, "elementwise") | ||
| 609 | + | ||
| 610 | + def test_standalone_full_is_not_elementwise_workload(self): | ||
| 611 | + def fn(x): | ||
| 612 | + return torch.full_like(x, 2.0) | ||
| 613 | + | ||
| 614 | + inputs = (torch.randn((257, 1031), device=device),) | ||
| 615 | + next_inputs = (torch.randn((263, 1031), device=device),) | ||
| 616 | + self._run_and_check(fn, inputs, next_inputs, None) | ||
| 617 | + | ||
| 618 | + def test_strided_reindex_is_not_elementwise_workload(self): | ||
| 619 | + def fn(x, y): | ||
| 620 | + return torch.relu(x[:, ::2] + y) | ||
| 621 | + | ||
| 622 | + inputs = ( | ||
| 623 | + torch.randn((257, 2062), device=device), | ||
| 624 | + torch.randn((257, 1031), device=device), | ||
| 625 | + ) | ||
| 626 | + next_inputs = ( | ||
| 627 | + torch.randn((263, 2062), device=device), | ||
| 628 | + torch.randn((263, 1031), device=device), | ||
| 629 | + ) | ||
| 630 | + self._run_and_check(fn, inputs, next_inputs, None) | ||
| 631 | + | ||
| 531 | # ============================================================ | 632 | # ============================================================ |
| 532 | # Part 2: Nonzero Data-Dependent Dynamic Shape Tests | 633 | # Part 2: Nonzero Data-Dependent Dynamic Shape Tests |
| 533 | # ============================================================ | 634 | # ============================================================ |
| @@ -752,6 +853,7 @@ class TestNonzeroDynamicShapes(DynamicShapeTestMixin, TestCase): | |||
| 752 | 853 | ||
| 753 | 854 | ||
| 754 | instantiate_parametrized_tests(TestFusionDynamicShapes) | 855 | instantiate_parametrized_tests(TestFusionDynamicShapes) |
| 856 | +instantiate_parametrized_tests(TestSymbolicGroupElementwise) | ||
| 755 | instantiate_parametrized_tests(TestNonzeroDynamicShapes) | 857 | instantiate_parametrized_tests(TestNonzeroDynamicShapes) |
| 756 | 858 | ||
| 757 | 859 | ||
| @@ -3,6 +3,7 @@ import sympy as sympy | |||
| 3 | import torch | 3 | import torch |
| 4 | from torch._inductor.codegen.simd import EnableReduction, DisableReduction | 4 | from torch._inductor.codegen.simd import EnableReduction, DisableReduction |
| 5 | from torch._inductor.codegen.triton import TritonKernel | 5 | from torch._inductor.codegen.triton import TritonKernel |
| 6 | +from torch._inductor.dependencies import MemoryDep | ||
| 6 | from torch._inductor.loop_body import MemoryUsageType | 7 | from torch._inductor.loop_body import MemoryUsageType |
| 7 | from torch._inductor.runtime.runtime_utils import next_power_of_2 | 8 | from torch._inductor.runtime.runtime_utils import next_power_of_2 |
| 8 | from torch._inductor.utils import ModularIndexing, sympy_subs | 9 | from torch._inductor.utils import ModularIndexing, sympy_subs |
| @@ -15,6 +16,10 @@ from ..config import num_vector_core, log | |||
| 15 | from ..runtime.symbolic_grouping import GroupFeatureSpec, GroupedKernelMeta | 16 | from ..runtime.symbolic_grouping import GroupFeatureSpec, GroupedKernelMeta |
| 16 | 17 | ||
| 17 | 18 | ||
| 19 | +_ELEMENTWISE_UNSUPPORTED_OPS = ("masked", "scan", "sort", "rand", "randn", "load_seed") | ||
| 20 | +_NEUTRAL_CONSTANT_OPS = frozenset(("constant", "store", "output")) | ||
| 21 | + | ||
| 22 | + | ||
| 18 | # split and tiling axis selector | 23 | # split and tiling axis selector |
| 19 | class SplitTiling: | 24 | class SplitTiling: |
| 20 | def __init__(self, kernel: TritonKernel): | 25 | def __init__(self, kernel: TritonKernel): |
| @@ -337,7 +342,7 @@ class SplitTiling: | |||
| 337 | _REDUCTION_BUCKETS = (8192,) | 342 | _REDUCTION_BUCKETS = (8192,) |
| 338 | _OUTER_BUCKETS = (256,) | 343 | _OUTER_BUCKETS = (256,) |
| 339 | 344 | ||
| 340 | - def _build_group_features(self, primary_axis): | 345 | + def _build_group_features(self, workload, primary_axis): |
| 341 | if self.kernel.persistent_reduction or self.kernel.inside_reduction: | 346 | if self.kernel.persistent_reduction or self.kernel.inside_reduction: |
| 342 | outer_names = self.non_reduction_axis_names() | 347 | outer_names = self.non_reduction_axis_names() |
| 343 | reduction_names = self.reduction_axis_names() | 348 | reduction_names = self.reduction_axis_names() |
| @@ -362,6 +367,15 @@ class SplitTiling: | |||
| 362 | ) | 367 | ) |
| 363 | ) | 368 | ) |
| 364 | return tuple(features) | 369 | return tuple(features) |
| 370 | + if workload == "elementwise": | ||
| 371 | + return ( | ||
| 372 | + GroupFeatureSpec( | ||
| 373 | + "elementwise_numel", | ||
| 374 | + "outer_product", | ||
| 375 | + self.all_axis_names(), | ||
| 376 | + (num_vector_core * 4096,), | ||
| 377 | + ), | ||
| 378 | + ) | ||
| 365 | return ( | 379 | return ( |
| 366 | GroupFeatureSpec( | 380 | GroupFeatureSpec( |
| 367 | "pointwise", | 381 | "pointwise", |
| @@ -371,8 +385,159 @@ class SplitTiling: | |||
| 371 | ), | 385 | ), |
| 372 | ) | 386 | ) |
| 373 | 387 | ||
| 388 | + @staticmethod | ||
| 389 | + def _alpha_rename_access_vars(dep): | ||
| 390 | + replacements = { | ||
| 391 | + var: sympy.Symbol(f"elementwise_dim_{idx}", integer=True, nonnegative=True) | ||
| 392 | + for idx, var in enumerate(dep.var_names) | ||
| 393 | + } | ||
| 394 | + return replacements | ||
| 395 | + | ||
| 396 | + def _make_access_signature(self, dep): | ||
| 397 | + if dep.is_indirect() or dep.mode is not None: | ||
| 398 | + return None | ||
| 399 | + normalized = dep.normalize() | ||
| 400 | + replacements = self._alpha_rename_access_vars(normalized) | ||
| 401 | + index = sympy_subs( | ||
| 402 | + normalized.index - normalized.get_offset(), | ||
| 403 | + replacements, | ||
| 404 | + ) | ||
| 405 | + sizes = tuple(sympy_subs(size, replacements) for size in normalized.size) | ||
| 406 | + return ( | ||
| 407 | + V.graph.sizevars.simplify(index), | ||
| 408 | + tuple(V.graph.sizevars.simplify(size) for size in sizes), | ||
| 409 | + ) | ||
| 410 | + | ||
| 411 | + @staticmethod | ||
| 412 | + def _same_access_signature(left, right): | ||
| 413 | + left_index, left_sizes = left | ||
| 414 | + right_index, right_sizes = right | ||
| 415 | + if len(left_sizes) != len(right_sizes): | ||
| 416 | + return False | ||
| 417 | + sizevars = V.graph.sizevars | ||
| 418 | + return sizevars.statically_known_equals(left_index, right_index) and all( | ||
| 419 | + sizevars.statically_known_equals(left_size, right_size) | ||
| 420 | + for left_size, right_size in zip(left_sizes, right_sizes) | ||
| 421 | + ) | ||
| 422 | + | ||
| 423 | + def _all_access_signatures_equal(self, signatures): | ||
| 424 | + if not signatures: | ||
| 425 | + return False | ||
| 426 | + reference = signatures[0] | ||
| 427 | + return all( | ||
| 428 | + self._same_access_signature(reference, signature) | ||
| 429 | + for signature in signatures[1:] | ||
| 430 | + ) | ||
| 431 | + | ||
| 432 | + @staticmethod | ||
| 433 | + def _has_unsupported_elementwise_semantics(node): | ||
| 434 | + return any(node._body.has_op(op) for op in _ELEMENTWISE_UNSUPPORTED_OPS) | ||
| 435 | + | ||
| 436 | + @staticmethod | ||
| 437 | + def _is_constant_only_body(node): | ||
| 438 | + op_names = set(node._body.op_counts) | ||
| 439 | + return bool(node._body.op_counts.get("constant")) and op_names <= _NEUTRAL_CONSTANT_OPS | ||
| 440 | + | ||
| 441 | + def _classify_elementwise_node(self, node): | ||
| 442 | + if node.has_side_effects() or node.has_aliasing_or_mutation(): | ||
| 443 | + return None | ||
| 444 | + if self._has_unsupported_elementwise_semantics(node): | ||
| 445 | + return None | ||
| 446 | + | ||
| 447 | + reads = tuple(node.read_writes.reads) | ||
| 448 | + writes = tuple(node.read_writes.writes) | ||
| 449 | + if not writes or node.read_writes.index_exprs: | ||
| 450 | + return None | ||
| 451 | + if not all(isinstance(dep, MemoryDep) for dep in writes): | ||
| 452 | + return None | ||
| 453 | + | ||
| 454 | + write_signatures = tuple(self._make_access_signature(dep) for dep in writes) | ||
| 455 | + if any(signature is None for signature in write_signatures): | ||
| 456 | + return None | ||
| 457 | + if not self._all_access_signatures_equal(write_signatures): | ||
| 458 | + return None | ||
| 459 | + reference = write_signatures[0] | ||
| 460 | + | ||
| 461 | + if not reads: | ||
| 462 | + # Constant-only producers are neutral only when the kernel also has | ||
| 463 | + # a direct external tensor read. | ||
| 464 | + if self._is_constant_only_body(node): | ||
| 465 | + return "neutral_constant", reference | ||
| 466 | + return None | ||
| 467 | + | ||
| 468 | + if not all(isinstance(dep, MemoryDep) for dep in reads): | ||
| 469 | + return None | ||
| 470 | + read_signatures = tuple(self._make_access_signature(dep) for dep in reads) | ||
| 471 | + if any(signature is None for signature in read_signatures): | ||
| 472 | + return None | ||
| 473 | + if not all( | ||
| 474 | + self._same_access_signature(reference, signature) | ||
| 475 | + for signature in read_signatures | ||
| 476 | + ): | ||
| 477 | + return None | ||
| 478 | + return "direct", reference | ||
| 479 | + | ||
| 480 | + @staticmethod | ||
| 481 | + def _neutral_outputs_are_consumed(classifications): | ||
| 482 | + # A neutral node must represent an internal value, not an independent | ||
| 483 | + # generated output fused horizontally into the same kernel. | ||
| 484 | + consumed_names = { | ||
| 485 | + dep.name | ||
| 486 | + for node, _, _ in classifications | ||
| 487 | + for dep in node.read_writes.reads | ||
| 488 | + if isinstance(dep, MemoryDep) | ||
| 489 | + } | ||
| 490 | + return all( | ||
| 491 | + dep.name in consumed_names | ||
| 492 | + for node, kind, _ in classifications | ||
| 493 | + if kind == "neutral_constant" | ||
| 494 | + for dep in node.read_writes.writes | ||
| 495 | + ) | ||
| 496 | + | ||
| 497 | + @staticmethod | ||
| 498 | + def _has_external_direct_read(classifications): | ||
| 499 | + produced_names = { | ||
| 500 | + dep.name | ||
| 501 | + for node, _, _ in classifications | ||
| 502 | + for dep in node.read_writes.writes | ||
| 503 | + if isinstance(dep, MemoryDep) | ||
| 504 | + } | ||
| 505 | + return any( | ||
| 506 | + dep.name not in produced_names | ||
| 507 | + for node, kind, _ in classifications | ||
| 508 | + if kind == "direct" | ||
| 509 | + for dep in node.read_writes.reads | ||
| 510 | + if isinstance(dep, MemoryDep) | ||
| 511 | + ) | ||
| 512 | + | ||
| 513 | + def _classify_group_workload(self, template): | ||
| 514 | + if template != "pointwise": | ||
| 515 | + return None | ||
| 516 | + nodes = tuple(self.kernel.features.scheduler_nodes()) | ||
| 517 | + if not nodes: | ||
| 518 | + return None | ||
| 519 | + | ||
| 520 | + classifications = [] | ||
| 521 | + for node in nodes: | ||
| 522 | + result = self._classify_elementwise_node(node) | ||
| 523 | + if result is None: | ||
| 524 | + return None | ||
| 525 | + kind, reference = result | ||
| 526 | + classifications.append((node, kind, reference)) | ||
| 527 | + | ||
| 528 | + references = tuple(reference for _, _, reference in classifications) | ||
| 529 | + if not self._all_access_signatures_equal(references): | ||
| 530 | + return None | ||
| 531 | + if not self._neutral_outputs_are_consumed(classifications): | ||
| 532 | + return None | ||
| 533 | + if not self._has_external_direct_read(classifications): | ||
| 534 | + return None | ||
| 535 | + return "elementwise" | ||
| 536 | + | ||
| 374 | def _build_grouped_meta(self): | 537 | def _build_grouped_meta(self): |
| 375 | dynamic_split_axes, static_split_axes = self._classify_split_axes() | 538 | dynamic_split_axes, static_split_axes = self._classify_split_axes() |
| 539 | + template = self._grouped_template_name() | ||
| 540 | + workload = self._classify_group_workload(template) | ||
| 376 | if dynamic_split_axes: | 541 | if dynamic_split_axes: |
| 377 | primary_axis = self._select_primary_group_axis(dynamic_split_axes) | 542 | primary_axis = self._select_primary_group_axis(dynamic_split_axes) |
| 378 | if primary_axis is None: | 543 | if primary_axis is None: |
| @@ -398,12 +563,13 @@ class SplitTiling: | |||
| 398 | runtime_block_arg_names = tuple( | 563 | runtime_block_arg_names = tuple( |
| 399 | f"{axis.name.upper()}BLOCK" for axis in self.kernel.split_axis | 564 | f"{axis.name.upper()}BLOCK" for axis in self.kernel.split_axis |
| 400 | ) | 565 | ) |
| 401 | - feature_specs = self._build_group_features(primary_axis) | 566 | + feature_specs = self._build_group_features(workload, primary_axis) |
| 402 | if not feature_specs: | 567 | if not feature_specs: |
| 403 | return None | 568 | return None |
| 404 | return GroupedKernelMeta( | 569 | return GroupedKernelMeta( |
| 405 | enabled=True, | 570 | enabled=True, |
| 406 | - template=self._grouped_template_name(), | 571 | + template=template, |
| 572 | + workload=workload, | ||
| 407 | primary_group_axis=primary_axis.name, | 573 | primary_group_axis=primary_axis.name, |
| 408 | static_split_axes=static_names, | 574 | static_split_axes=static_names, |
| 409 | secondary_runtime_symbolic_axes=secondary_names, | 575 | secondary_runtime_symbolic_axes=secondary_names, |
| @@ -1990,6 +1990,7 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 1990 | { | 1990 | { |
| 1991 | "group_enabled": bool(grouped_meta.enabled), | 1991 | "group_enabled": bool(grouped_meta.enabled), |
| 1992 | "group_template": grouped_meta.template, | 1992 | "group_template": grouped_meta.template, |
| 1993 | + "group_workload": grouped_meta.workload, | ||
| 1993 | "primary_group_axis": grouped_meta.primary_group_axis, | 1994 | "primary_group_axis": grouped_meta.primary_group_axis, |
| 1994 | "static_split_axes": grouped_meta.static_split_axes, | 1995 | "static_split_axes": grouped_meta.static_split_axes, |
| 1995 | "secondary_runtime_symbolic_axes": grouped_meta.secondary_runtime_symbolic_axes, | 1996 | "secondary_runtime_symbolic_axes": grouped_meta.secondary_runtime_symbolic_axes, |
| @@ -2004,6 +2005,7 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2004 | { | 2005 | { |
| 2005 | "group_enabled": False, | 2006 | "group_enabled": False, |
| 2006 | "group_template": None, | 2007 | "group_template": None, |
| 2008 | + "group_workload": None, | ||
| 2007 | "primary_group_axis": None, | 2009 | "primary_group_axis": None, |
| 2008 | "static_split_axes": (), | 2010 | "static_split_axes": (), |
| 2009 | "secondary_runtime_symbolic_axes": (), | 2011 | "secondary_runtime_symbolic_axes": (), |
| @@ -2096,6 +2098,7 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2096 | ) | 2098 | ) |
| 2097 | inductor_meta["group_enabled"] = False | 2099 | inductor_meta["group_enabled"] = False |
| 2098 | inductor_meta["group_template"] = None | 2100 | inductor_meta["group_template"] = None |
| 2101 | + inductor_meta["group_workload"] = None | ||
| 2099 | inductor_meta["primary_group_axis"] = None | 2102 | inductor_meta["primary_group_axis"] = None |
| 2100 | inductor_meta["static_split_axes"] = () | 2103 | inductor_meta["static_split_axes"] = () |
| 2101 | inductor_meta["secondary_runtime_symbolic_axes"] = () | 2104 | inductor_meta["secondary_runtime_symbolic_axes"] = () |
| @@ -51,6 +51,7 @@ class UnsupportedGroupedPlan(RuntimeError): | |||
| 51 | class GroupedKernelMeta: | 51 | class GroupedKernelMeta: |
| 52 | enabled: bool | 52 | enabled: bool |
| 53 | template: str | 53 | template: str |
| 54 | + workload: str | None | ||
| 54 | primary_group_axis: str | None | 55 | primary_group_axis: str | None |
| 55 | static_split_axes: tuple[str, ...] | 56 | static_split_axes: tuple[str, ...] |
| 56 | secondary_runtime_symbolic_axes: tuple[str, ...] | 57 | secondary_runtime_symbolic_axes: tuple[str, ...] |
| @@ -61,6 +62,7 @@ class GroupedKernelMeta: | |||
| 61 | return { | 62 | return { |
| 62 | "enabled": self.enabled, | 63 | "enabled": self.enabled, |
| 63 | "template": self.template, | 64 | "template": self.template, |
| 65 | + "workload": self.workload, | ||
| 64 | "primary_group_axis": self.primary_group_axis, | 66 | "primary_group_axis": self.primary_group_axis, |
| 65 | "static_split_axes": list(self.static_split_axes), | 67 | "static_split_axes": list(self.static_split_axes), |
| 66 | "secondary_runtime_symbolic_axes": list( | 68 | "secondary_runtime_symbolic_axes": list( |
| @@ -83,6 +85,7 @@ class GroupedKernelMeta: | |||
| 83 | return GroupedKernelMeta( | 85 | return GroupedKernelMeta( |
| 84 | enabled=_require_bool(payload.get("enabled"), "enabled"), | 86 | enabled=_require_bool(payload.get("enabled"), "enabled"), |
| 85 | template=_require_str(payload.get("template"), "template"), | 87 | template=_require_str(payload.get("template"), "template"), |
| 88 | + workload=_require_group_workload(payload.get("workload")), | ||
| 86 | primary_group_axis=_require_optional_str( | 89 | primary_group_axis=_require_optional_str( |
| 87 | payload.get("primary_group_axis"), "primary_group_axis" | 90 | payload.get("primary_group_axis"), "primary_group_axis" |
| 88 | ), | 91 | ), |
| @@ -384,6 +387,13 @@ def _require_optional_str(value: object, field_name: str) -> str | None: | |||
| 384 | return _require_str(value, field_name) | 387 | return _require_str(value, field_name) |
| 385 | 388 | ||
| 386 | 389 | ||
| 390 | +def _require_group_workload(value: object) -> str | None: | ||
| 391 | + workload = _require_optional_str(value, "workload") | ||
| 392 | + if workload not in (None, "elementwise"): | ||
| 393 | + raise ValueError(f"unsupported group workload: {workload}") | ||
| 394 | + return workload | ||
| 395 | + | ||
| 396 | + | ||
| 387 | def _require_bool(value: object, field_name: str) -> bool: | 397 | def _require_bool(value: object, field_name: str) -> bool: |
| 388 | if not isinstance(value, bool): | 398 | if not isinstance(value, bool): |
| 389 | raise TypeError(f"{field_name} must be a bool, got {type(value).__name__}") | 399 | raise TypeError(f"{field_name} must be a bool, got {type(value).__name__}") |