(
conv_type,
conv_kind,
tile_descriptions,
src_layout,
flt_layout,
dst_layout,
dst_type,
min_cc,
src_align=32,
flt_align=32,
dst_align=32,
use_special_optimization=SpecialOptimizeDesc.NoneSpecialOpt,
implicit_gemm_mode=ImplicitGemmMode.GemmNT,
without_shared_load=False,
required_cuda_ver_major=9,
required_cuda_ver_minor=2,
)
| 639 | |
| 640 | # |
| 641 | def GenerateConv2d( |
| 642 | conv_type, |
| 643 | conv_kind, |
| 644 | tile_descriptions, |
| 645 | src_layout, |
| 646 | flt_layout, |
| 647 | dst_layout, |
| 648 | dst_type, |
| 649 | min_cc, |
| 650 | src_align=32, |
| 651 | flt_align=32, |
| 652 | dst_align=32, |
| 653 | use_special_optimization=SpecialOptimizeDesc.NoneSpecialOpt, |
| 654 | implicit_gemm_mode=ImplicitGemmMode.GemmNT, |
| 655 | without_shared_load=False, |
| 656 | required_cuda_ver_major=9, |
| 657 | required_cuda_ver_minor=2, |
| 658 | ): |
| 659 | operations = [] |
| 660 | |
| 661 | element_epilogue = DataType.f32 |
| 662 | if ( |
| 663 | conv_type == ConvType.DepthwiseConvolution |
| 664 | or conv_type == ConvType.RegionRestrictedConvolution |
| 665 | ): |
| 666 | if conv_kind == ConvKind.Fprop: |
| 667 | swizzling_functor = SwizzlingFunctor.DepthwiseConvolutionFprop |
| 668 | elif conv_kind == ConvKind.Dgrad: |
| 669 | swizzling_functor = SwizzlingFunctor.DepthwiseConvolutionDgrad |
| 670 | else: |
| 671 | assert conv_kind == ConvKind.Wgrad |
| 672 | swizzling_functor = SwizzlingFunctor.DepthwiseConvolutionWgrad |
| 673 | elif conv_type == ConvType.Convolution: |
| 674 | if conv_kind == ConvKind.Fprop: |
| 675 | if implicit_gemm_mode == ImplicitGemmMode.GemmTN: |
| 676 | swizzling_functor = SwizzlingFunctor.ConvFpropTrans |
| 677 | else: |
| 678 | swizzling_functor = SwizzlingFunctor.ConvFpropNCxHWx |
| 679 | else: |
| 680 | if implicit_gemm_mode == ImplicitGemmMode.GemmTN: |
| 681 | swizzling_functor = SwizzlingFunctor.ConvDgradTrans |
| 682 | else: |
| 683 | swizzling_functor = SwizzlingFunctor.ConvDgradNCxHWx |
| 684 | |
| 685 | # skip rule |
| 686 | def filter_tile_with_layout(tile: TileDescription, layout: LayoutType) -> bool: |
| 687 | return ( |
| 688 | layout == LayoutType.TensorNC32HW32 and tile.threadblock_shape[0] % 32 != 0 |
| 689 | ) |
| 690 | |
| 691 | # rule for bias_type and epilogues |
| 692 | def get_bias_type_and_epilogues( |
| 693 | tile: TileDescription, out_dtype: DataType |
| 694 | ) -> Tuple[DataType, List[EpilogueFunctor]]: |
| 695 | if ( |
| 696 | tile.math_instruction.element_accumulator == DataType.s32 |
| 697 | and out_dtype != DataType.f32 |
| 698 | ): |
no test coverage detected