| 628 | |
| 629 | |
| 630 | class ExternKernelCaller(ChoiceCaller): |
| 631 | def __init__( |
| 632 | self, |
| 633 | choice: ExternKernelChoice, |
| 634 | input_nodes, |
| 635 | layout, |
| 636 | kwargs=None, |
| 637 | *, |
| 638 | has_out_variant=True, |
| 639 | ): |
| 640 | super().__init__(choice.name, input_nodes, layout) |
| 641 | self.choice = choice |
| 642 | self.kwargs = kwargs or {} |
| 643 | self.has_out_variant = has_out_variant |
| 644 | |
| 645 | def __str__(self): |
| 646 | return f"ExternKernelCaller({self.choice.call_name()})" |
| 647 | |
| 648 | def benchmark(self, *args, out): |
| 649 | if self.has_out_variant: |
| 650 | return super().benchmark(*args, out=out) |
| 651 | else: |
| 652 | algo = self.to_callable() |
| 653 | out_new = algo(*args) |
| 654 | torch._C._dynamo.guards.assert_size_stride( |
| 655 | out_new, tuple(out.size()), tuple(out.stride()) |
| 656 | ) |
| 657 | out.copy_(out_new) # for correctness checking |
| 658 | return do_bench(lambda: algo(*args)) |
| 659 | |
| 660 | def to_callable(self): |
| 661 | fn = self.choice.to_callable() |
| 662 | if self.kwargs: |
| 663 | return functools.partial(fn, **self.kwargs) |
| 664 | else: |
| 665 | return fn |
| 666 | |
| 667 | def hash_key(self): |
| 668 | return "-".join( |
| 669 | [ |
| 670 | self.choice.name, |
| 671 | *[ |
| 672 | f"{kwarg}={repr(self.kwargs[kwarg])}" |
| 673 | for kwarg in sorted(self.kwargs.keys()) |
| 674 | ], |
| 675 | self.choice.hash_key(), |
| 676 | ] |
| 677 | ) |
| 678 | |
| 679 | def output_node(self): |
| 680 | cls: Union[Type[ir.ExternKernelOut], Type[ir.ExternKernelAlloc]] |
| 681 | if self.has_out_variant: |
| 682 | cls = ir.ExternKernelOut |
| 683 | else: |
| 684 | cls = ir.ExternKernelAlloc |
| 685 | return ir.TensorBox.create( |
| 686 | cls( |
| 687 | layout=self.layout, |
no outgoing calls
no test coverage detected
searching dependent graphs…