MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / GenerateConv2d

Function GenerateConv2d

dnn/scripts/cutlass_generator/conv2d_operation.py:641–865  ·  view source on GitHub ↗
(
    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,
)

Source from the content-addressed store, hash-verified

639
640#
641def 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 ):

Calls 9

filter_tile_with_layoutFunction · 0.85
get_flt_alignFunction · 0.85
get_dst_alignFunction · 0.85
maxFunction · 0.85
Conv2dOperationClass · 0.85
TensorDescriptionClass · 0.70
appendMethod · 0.45

Tested by

no test coverage detected