OpenI
/
MegEngine
Not watched
1
0
Fork 0
Code
Releases 31 Wiki Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: c42ce93705
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'c42ce93705'
${ noResults }
MegEngine/dnn/scripts/cutlass_generator/conv2d_operation.py

974 lines
35 kB
Raw Blame History 

  1. #

  2. # \file generator.py

  3. #

  4. # \brief Generates the CUTLASS Library's instances

  5. #

  6. #

  7. 

  8. import enum

  9. import os.path

  10. import shutil

  11. from typing import Tuple, List

  12. 

  13. from library import *

  14. 

  15. ###################################################################################################

  16. 

  17. #

  18. class Conv2dOperation:

  19.     #

  20.     def __init__(

  21.         self,

  22.         conv_kind,

  23.         conv_type,

  24.         arch,

  25.         tile_description,

  26.         src,

  27.         flt,

  28.         bias,

  29.         dst,

  30.         element_epilogue,

  31.         epilogue_functor=EpilogueFunctor.LinearCombination,

  32.         swizzling_functor=SwizzlingFunctor.Identity4,

  33.         special_optimization=SpecialOptimizeDesc.NoneSpecialOpt,

  34.         implicit_gemm_mode=ImplicitGemmMode.GemmNT,

  35.         without_shared_load=False,

  36.         required_cuda_ver_major=9,

  37.         required_cuda_ver_minor=2,

  38.     ):

  39. 

  40.         self.operation_kind = OperationKind.Conv2d

  41.         self.conv_kind = conv_kind

  42.         self.arch = arch

  43.         self.tile_description = tile_description

  44.         self.conv_type = conv_type

  45.         self.src = src

  46.         self.flt = flt

  47.         self.bias = bias

  48.         self.dst = dst

  49.         self.element_epilogue = element_epilogue

  50.         self.epilogue_functor = epilogue_functor

  51.         self.swizzling_functor = swizzling_functor

  52.         self.special_optimization = special_optimization

  53.         self.implicit_gemm_mode = implicit_gemm_mode

  54.         self.without_shared_load = without_shared_load

  55.         self.required_cuda_ver_major = required_cuda_ver_major

  56.         self.required_cuda_ver_minor = required_cuda_ver_minor

  57. 

  58.     #

  59.     def accumulator_type(self):

  60.         accum = self.tile_description.math_instruction.element_accumulator

  61. 

  62.         return accum

  63. 

  64.     #

  65.     def core_name(self):

  66.         """ The basic operation kind is prefixed with a letter indicating the accumulation type. """

  67. 

  68.         intermediate_type = ""

  69. 

  70.         if self.tile_description.math_instruction.opcode_class == OpcodeClass.TensorOp:

  71.             inst_shape = "%d%d%d" % tuple(

  72.                 self.tile_description.math_instruction.instruction_shape

  73.             )

  74.             if (

  75.                 self.tile_description.math_instruction.element_a != self.flt.element

  76.                 and self.tile_description.math_instruction.element_a

  77.                 != self.accumulator_type()

  78.             ):

  79.                 intermediate_type = DataTypeNames[

  80.                     self.tile_description.math_instruction.element_a

  81.                 ]

  82.         else:

  83.             inst_shape = ""

  84. 

  85.         special_opt = ""

  86.         if self.special_optimization == SpecialOptimizeDesc.ConvFilterUnity:

  87.             special_opt = "_1x1"

  88.         elif self.special_optimization == SpecialOptimizeDesc.DeconvDoubleUpsampling:

  89.             special_opt = "_s2"

  90. 

  91.         reorder_k = ""

  92.         if self.without_shared_load:

  93.             reorder_k = "_roc"

  94. 

  95.         conv_type_name = ""

  96.         if self.conv_type == ConvType.DepthwiseConvolution:

  97.             conv_type_name = "dw"

  98. 

  99.         return "%s%s%s%s%s%s%s_%s" % (

  100.             ShortDataTypeNames[self.accumulator_type()],

  101.             inst_shape,

  102.             intermediate_type,

  103.             conv_type_name,

  104.             ConvKindNames[self.conv_kind],

  105.             special_opt,

  106.             reorder_k,

  107.             ShortEpilogueNames[self.epilogue_functor],

  108.         )

  109. 

  110.     #

  111.     def extended_name(self):

  112.         if (

  113.             self.dst.element

  114.             != self.tile_description.math_instruction.element_accumulator

  115.         ):

  116.             if self.src.element != self.flt.element:

  117.                 extended_name = (

  118.                     "${element_dst}_${core_name}_${element_src}_${element_flt}"

  119.                 )

  120.             elif self.src.element == self.flt.element:

  121.                 extended_name = "${element_dst}_${core_name}_${element_src}"

  122.         else:

  123.             if self.src.element != self.flt.element:

  124.                 extended_name = "${core_name}_${element_src}_${element_flt}"

  125.             elif self.src.element == self.flt.element:

  126.                 extended_name = "${core_name}_${element_src}"

  127. 

  128.         extended_name = SubstituteTemplate(

  129.             extended_name,

  130.             {

  131.                 "element_src": DataTypeNames[self.src.element],

  132.                 "element_flt": DataTypeNames[self.flt.element],

  133.                 "element_dst": DataTypeNames[self.dst.element],

  134.                 "core_name": self.core_name(),

  135.             },

  136.         )

  137. 

  138.         return extended_name

  139. 

  140.     #

  141.     def layout_name(self):

  142.         if self.src.layout == self.dst.layout:

  143.             layout_name = "${src_layout}_${flt_layout}"

  144.         else:

  145.             layout_name = "${src_layout}_${flt_layout}_${dst_layout}"

  146. 

  147.         layout_name = SubstituteTemplate(

  148.             layout_name,

  149.             {

  150.                 "src_layout": ShortLayoutTypeNames[self.src.layout],

  151.                 "flt_layout": ShortLayoutTypeNames[self.flt.layout],

  152.                 "dst_layout": ShortLayoutTypeNames[self.dst.layout],

  153.             },

  154.         )

  155. 

  156.         return layout_name

  157. 

  158.     #

  159.     def configuration_name(self):

  160.         """ The full procedural name indicates architecture, extended name, tile size, and layout. """

  161. 

  162.         opcode_class_name = OpcodeClassNames[

  163.             self.tile_description.math_instruction.opcode_class

  164.         ]

  165. 

  166.         warp_shape = [

  167.             int(

  168.                 self.tile_description.threadblock_shape[idx]

  169.                 / self.tile_description.warp_count[idx]

  170.             )

  171.             for idx in range(3)

  172.         ]

  173. 

  174.         threadblock = "%dx%dx%d_%dx%dx%d_%d" % (

  175.             self.tile_description.threadblock_shape[0],

  176.             self.tile_description.threadblock_shape[1],

  177.             self.tile_description.threadblock_shape[2],

  178.             warp_shape[0],

  179.             warp_shape[1],

  180.             warp_shape[2],

  181.             self.tile_description.stages,

  182.         )

  183. 

  184.         alignment = "align%dx%d" % (self.src.alignment, self.flt.alignment)

  185. 

  186.         configuration_name = "cutlass_${opcode_class}_${extended_name}_${threadblock}_${layout}_${alignment}"

  187. 

  188.         return SubstituteTemplate(

  189.             configuration_name,

  190.             {

  191.                 "opcode_class": opcode_class_name,

  192.                 "extended_name": self.extended_name(),

  193.                 "threadblock": threadblock,

  194.                 "layout": self.layout_name(),

  195.                 "alignment": alignment,

  196.             },

  197.         )

  198. 

  199.     #

  200.     def procedural_name(self):

  201.         """ The full procedural name indicates architecture, extended name, tile size, and layout. """

  202.         return self.configuration_name()

  203. 

  204. 

  205. ###################################################################################################

  206. #

  207. # Emits single instances of a CUTLASS device-wide operator

  208. #

  209. ###################################################################################################

  210. 

  211. 

  212. class EmitConv2dInstance:

  213.     def __init__(self):

  214.         self.template = """

  215. // kernel instance "${operation_name}" generated by cutlass generator

  216. using Convolution = 

  217.   typename cutlass::conv::device::Convolution<

  218.     ${element_src}, 

  219.     ${layout_src},

  220.     ${element_flt}, 

  221.     ${layout_flt},

  222.     ${element_dst}, 

  223.     ${layout_dst},

  224.     ${element_bias}, 

  225.     ${layout_bias}, 

  226.     ${element_accumulator}, 

  227.     ${conv_type},

  228.     ${opcode_class},

  229.     ${arch},

  230.     cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,

  231.     cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k}>,

  232.     cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,

  233.     ${epilogue_functor}<

  234.       ${element_dst},

  235.       ${epilogue_vector_length},

  236.       ${element_accumulator}, 

  237.       ${element_bias}, 

  238.       ${element_epilogue}

  239.     >,

  240.     ${swizzling_functor},     

  241.     ${stages},

  242.     ${alignment_src}, 

  243.     ${alignment_filter}, 

  244.     ${special_optimization}, 

  245.     ${math_operator},

  246.     ${implicit_gemm_mode}, 

  247.     ${without_shared_load}>;

  248. """

  249. 

  250.     def emit(self, operation):

  251. 

  252.         warp_shape = [

  253.             int(

  254.                 operation.tile_description.threadblock_shape[idx]

  255.                 / operation.tile_description.warp_count[idx]

  256.             )

  257.             for idx in range(3)

  258.         ]

  259. 

  260.         epilogue_vector_length = int(

  261.             min(operation.dst.alignment * DataTypeSize[operation.dst.element], 128)

  262.             / DataTypeSize[operation.dst.element]

  263.         )

  264. 

  265.         values = {

  266.             "operation_name": operation.procedural_name(),

  267.             "conv_type": ConvTypeTag[operation.conv_type],

  268.             "element_src": DataTypeTag[operation.src.element],

  269.             "layout_src": LayoutTag[operation.src.layout],

  270.             "element_flt": DataTypeTag[operation.flt.element],

  271.             "layout_flt": LayoutTag[operation.flt.layout],

  272.             "element_dst": DataTypeTag[operation.dst.element],

  273.             "layout_dst": LayoutTag[operation.dst.layout],

  274.             "element_bias": DataTypeTag[operation.bias.element],

  275.             "layout_bias": LayoutTag[operation.bias.layout],

  276.             "element_accumulator": DataTypeTag[operation.accumulator_type()],

  277.             "opcode_class": OpcodeClassTag[

  278.                 operation.tile_description.math_instruction.opcode_class

  279.             ],

  280.             "arch": "cutlass::arch::Sm%d" % operation.arch,

  281.             "threadblock_shape_m": str(operation.tile_description.threadblock_shape[0]),

  282.             "threadblock_shape_n": str(operation.tile_description.threadblock_shape[1]),

  283.             "threadblock_shape_k": str(operation.tile_description.threadblock_shape[2]),

  284.             "warp_shape_m": str(warp_shape[0]),

  285.             "warp_shape_n": str(warp_shape[1]),

  286.             "warp_shape_k": str(warp_shape[2]),

  287.             "instruction_shape_m": str(

  288.                 operation.tile_description.math_instruction.instruction_shape[0]

  289.             ),

  290.             "instruction_shape_n": str(

  291.                 operation.tile_description.math_instruction.instruction_shape[1]

  292.             ),

  293.             "instruction_shape_k": str(

  294.                 operation.tile_description.math_instruction.instruction_shape[2]

  295.             ),

  296.             "epilogue_vector_length": str(epilogue_vector_length),

  297.             "epilogue_functor": EpilogueFunctorTag[operation.epilogue_functor],

  298.             "element_epilogue": str(DataTypeTag[operation.element_epilogue]),

  299.             "swizzling_functor": SwizzlingFunctorTag[operation.swizzling_functor],

  300.             "stages": str(operation.tile_description.stages),

  301.             "alignment_src": str(operation.src.alignment),

  302.             "alignment_filter": str(operation.flt.alignment),

  303.             "special_optimization": SpecialOptimizeDescTag[

  304.                 operation.special_optimization

  305.             ],

  306.             "math_operator": MathOperationTag[

  307.                 operation.tile_description.math_instruction.math_operation

  308.             ],

  309.             "implicit_gemm_mode": ImplicitGemmModeTag[operation.implicit_gemm_mode],

  310.             "without_shared_load": str(operation.without_shared_load).lower(),

  311.         }

  312. 

  313.         return SubstituteTemplate(self.template, values)

  314. 

  315. 

  316. class EmitDeconvInstance:

  317.     def __init__(self):

  318.         self.template = """

  319. // kernel instance "${operation_name}" generated by cutlass generator

  320. using Convolution = 

  321.   typename cutlass::conv::device::Deconvolution<

  322.     ${element_src}, 

  323.     ${layout_src},

  324.     ${element_flt}, 

  325.     ${layout_flt},

  326.     ${element_dst}, 

  327.     ${layout_dst},

  328.     ${element_bias}, 

  329.     ${layout_bias}, 

  330.     ${element_accumulator}, 

  331.     ${conv_type},

  332.     ${opcode_class},

  333.     ${arch},

  334.     cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,

  335.     cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k}>,

  336.     cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,

  337.     ${epilogue_functor}<

  338.       ${element_dst},

  339.       ${epilogue_vector_length},

  340.       ${element_accumulator}, 

  341.       ${element_bias}, 

  342.       ${element_epilogue}

  343.     >,

  344.     ${swizzling_functor},     

  345.     ${stages},

  346.     ${alignment_src}, 

  347.     ${alignment_filter}, 

  348.     ${special_optimization}, 

  349.     ${math_operator},

  350.     ${implicit_gemm_mode}>;

  351. """

  352. 

  353.     def emit(self, operation):

  354. 

  355.         warp_shape = [

  356.             int(

  357.                 operation.tile_description.threadblock_shape[idx]

  358.                 / operation.tile_description.warp_count[idx]

  359.             )

  360.             for idx in range(3)

  361.         ]

  362. 

  363.         epilogue_vector_length = int(

  364.             min(operation.dst.alignment * DataTypeSize[operation.dst.element], 128)

  365.             / DataTypeSize[operation.dst.element]

  366.         )

  367. 

  368.         values = {

  369.             "operation_name": operation.procedural_name(),

  370.             "conv_type": ConvTypeTag[operation.conv_type],

  371.             "element_src": DataTypeTag[operation.src.element],

  372.             "layout_src": LayoutTag[operation.src.layout],

  373.             "element_flt": DataTypeTag[operation.flt.element],

  374.             "layout_flt": LayoutTag[operation.flt.layout],

  375.             "element_dst": DataTypeTag[operation.dst.element],

  376.             "layout_dst": LayoutTag[operation.dst.layout],

  377.             "element_bias": DataTypeTag[operation.bias.element],

  378.             "layout_bias": LayoutTag[operation.bias.layout],

  379.             "element_accumulator": DataTypeTag[operation.accumulator_type()],

  380.             "opcode_class": OpcodeClassTag[

  381.                 operation.tile_description.math_instruction.opcode_class

  382.             ],

  383.             "arch": "cutlass::arch::Sm%d" % operation.arch,

  384.             "threadblock_shape_m": str(operation.tile_description.threadblock_shape[0]),

  385.             "threadblock_shape_n": str(operation.tile_description.threadblock_shape[1]),

  386.             "threadblock_shape_k": str(operation.tile_description.threadblock_shape[2]),

  387.             "warp_shape_m": str(warp_shape[0]),

  388.             "warp_shape_n": str(warp_shape[1]),

  389.             "warp_shape_k": str(warp_shape[2]),

  390.             "instruction_shape_m": str(

  391.                 operation.tile_description.math_instruction.instruction_shape[0]

  392.             ),

  393.             "instruction_shape_n": str(

  394.                 operation.tile_description.math_instruction.instruction_shape[1]

  395.             ),

  396.             "instruction_shape_k": str(

  397.                 operation.tile_description.math_instruction.instruction_shape[2]

  398.             ),

  399.             "epilogue_vector_length": str(epilogue_vector_length),

  400.             "epilogue_functor": EpilogueFunctorTag[operation.epilogue_functor],

  401.             "element_epilogue": str(DataTypeTag[operation.element_epilogue]),

  402.             "swizzling_functor": SwizzlingFunctorTag[operation.swizzling_functor],

  403.             "stages": str(operation.tile_description.stages),

  404.             "alignment_src": str(operation.src.alignment),

  405.             "alignment_filter": str(operation.flt.alignment),

  406.             "special_optimization": SpecialOptimizeDescTag[

  407.                 operation.special_optimization

  408.             ],

  409.             "math_operator": MathOperationTag[

  410.                 operation.tile_description.math_instruction.math_operation

  411.             ],

  412.             "implicit_gemm_mode": ImplicitGemmModeTag[operation.implicit_gemm_mode],

  413.         }

  414. 

  415.         return SubstituteTemplate(self.template, values)

  416. 

  417. 

  418. class EmitConvolutionBackwardFilterInstance:

  419.     def __init__(self):

  420.         self.template = """

  421. // kernel instance "${operation_name}" generated by cutlass generator

  422. using Convolution = 

  423.   typename cutlass::conv::device::ConvolutionBackwardFilter<

  424.     ${element_src}, 

  425.     ${layout_src},

  426.     ${element_diff}, 

  427.     ${layout_diff},

  428.     ${element_grad}, 

  429.     ${layout_grad},

  430.     ${element_accumulator}, 

  431.     ${conv_type},

  432.     ${opcode_class},

  433.     ${arch},

  434.     cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,

  435.     cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k}>,

  436.     cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,

  437.     ${epilogue_functor}<

  438.       ${element_grad},

  439.       ${epilogue_vector_length},

  440.       ${element_accumulator}, 

  441.       ${element_epilogue}

  442.     >,

  443.     ${swizzling_functor},     

  444.     ${stages},

  445.     ${alignment_src}, 

  446.     ${alignment_diff}, 

  447.     ${special_optimization}, 

  448.     ${math_operator},

  449.     ${implicit_gemm_mode}>;

  450. """

  451. 

  452.     def emit(self, operation):

  453. 

  454.         warp_shape = [

  455.             int(

  456.                 operation.tile_description.threadblock_shape[idx]

  457.                 / operation.tile_description.warp_count[idx]

  458.             )

  459.             for idx in range(3)

  460.         ]

  461. 

  462.         epilogue_vector_length = int(

  463.             min(operation.dst.alignment * DataTypeSize[operation.dst.element], 128)

  464.             / DataTypeSize[operation.dst.element]

  465.         )

  466. 

  467.         values = {

  468.             "operation_name": operation.procedural_name(),

  469.             "conv_type": ConvTypeTag[operation.conv_type],

  470.             "element_src": DataTypeTag[operation.src.element],

  471.             "layout_src": LayoutTag[operation.src.layout],

  472.             "element_diff": DataTypeTag[operation.flt.element],

  473.             "layout_diff": LayoutTag[operation.flt.layout],

  474.             "element_grad": DataTypeTag[operation.dst.element],

  475.             "layout_grad": LayoutTag[operation.dst.layout],

  476.             "element_accumulator": DataTypeTag[operation.accumulator_type()],

  477.             "opcode_class": OpcodeClassTag[

  478.                 operation.tile_description.math_instruction.opcode_class

  479.             ],

  480.             "arch": "cutlass::arch::Sm%d" % operation.arch,

  481.             "threadblock_shape_m": str(operation.tile_description.threadblock_shape[0]),

  482.             "threadblock_shape_n": str(operation.tile_description.threadblock_shape[1]),

  483.             "threadblock_shape_k": str(operation.tile_description.threadblock_shape[2]),

  484.             "warp_shape_m": str(warp_shape[0]),

  485.             "warp_shape_n": str(warp_shape[1]),

  486.             "warp_shape_k": str(warp_shape[2]),

  487.             "instruction_shape_m": str(

  488.                 operation.tile_description.math_instruction.instruction_shape[0]

  489.             ),

  490.             "instruction_shape_n": str(

  491.                 operation.tile_description.math_instruction.instruction_shape[1]

  492.             ),

  493.             "instruction_shape_k": str(

  494.                 operation.tile_description.math_instruction.instruction_shape[2]

  495.             ),

  496.             "epilogue_vector_length": str(epilogue_vector_length),

  497.             "epilogue_functor": EpilogueFunctorTag[operation.epilogue_functor],

  498.             "element_epilogue": str(DataTypeTag[operation.element_epilogue]),

  499.             "swizzling_functor": SwizzlingFunctorTag[operation.swizzling_functor],

  500.             "stages": str(operation.tile_description.stages),

  501.             "alignment_src": str(operation.src.alignment),

  502.             "alignment_diff": str(operation.flt.alignment),

  503.             "special_optimization": SpecialOptimizeDescTag[

  504.                 operation.special_optimization

  505.             ],

  506.             "math_operator": MathOperationTag[

  507.                 operation.tile_description.math_instruction.math_operation

  508.             ],

  509.             "implicit_gemm_mode": ImplicitGemmModeTag[operation.implicit_gemm_mode],

  510.         }

  511. 

  512.         return SubstituteTemplate(self.template, values)

  513. 

  514. 

  515. ###################################################################################################

  516. #

  517. # Generator functions for all layouts

  518. #

  519. ###################################################################################################

  520. 

  521. #

  522. def GenerateConv2d(

  523.     conv_type,

  524.     conv_kind,

  525.     tile_descriptions,

  526.     src_layout,

  527.     flt_layout,

  528.     dst_layout,

  529.     dst_type,

  530.     min_cc,

  531.     src_align=32,

  532.     flt_align=32,

  533.     dst_align=32,

  534.     use_special_optimization=SpecialOptimizeDesc.NoneSpecialOpt,

  535.     implicit_gemm_mode=ImplicitGemmMode.GemmNT,

  536.     without_shared_load=False,

  537.     required_cuda_ver_major=9,

  538.     required_cuda_ver_minor=2,

  539. ):

  540.     operations = []

  541. 

  542.     element_epilogue = DataType.f32

  543.     if conv_type == ConvType.DepthwiseConvolution:

  544.         if conv_kind == ConvKind.Fprop:

  545.             swizzling_functor = SwizzlingFunctor.DepthwiseConvolutionFprop

  546.         elif conv_kind == ConvKind.Dgrad:

  547.             swizzling_functor = SwizzlingFunctor.DepthwiseConvolutionDgrad

  548.         else:

  549.             assert conv_kind == ConvKind.Wgrad

  550.             swizzling_functor = SwizzlingFunctor.DepthwiseConvolutionWgrad

  551.     elif conv_type == ConvType.Convolution:

  552.         if conv_kind == ConvKind.Fprop:

  553.             if implicit_gemm_mode == ImplicitGemmMode.GemmTN:

  554.                 swizzling_functor = SwizzlingFunctor.ConvFpropTrans

  555.             else:

  556.                 swizzling_functor = SwizzlingFunctor.ConvFpropNCxHWx

  557.         else:

  558.             if implicit_gemm_mode == ImplicitGemmMode.GemmTN:

  559.                 swizzling_functor = SwizzlingFunctor.ConvDgradTrans

  560.             else:

  561.                 swizzling_functor = SwizzlingFunctor.ConvDgradNCxHWx

  562. 

  563.     # skip rule

  564.     def filter_tile_with_layout(tile: TileDescription, layout: LayoutType) -> bool:

  565.         return (

  566.             layout == LayoutType.TensorNC32HW32 and tile.threadblock_shape[0] % 32 != 0

  567.         )

  568. 

  569.     # rule for bias_type and epilogues

  570.     def get_bias_type_and_epilogues(

  571.         tile: TileDescription, out_dtype: DataType

  572.     ) -> Tuple[DataType, List[EpilogueFunctor]]:

  573.         if (

  574.             tile.math_instruction.element_accumulator == DataType.s32

  575.             and out_dtype != DataType.f32

  576.         ):

  577.             bias_type = DataType.s32

  578.             if tile.math_instruction.element_b == DataType.u4:

  579.                 epilogues = [

  580.                     EpilogueFunctor.BiasAddLinearCombinationClamp,

  581.                     EpilogueFunctor.BiasAddLinearCombinationReluClamp,

  582.                 ]

  583.             else:

  584.                 epilogues = [

  585.                     EpilogueFunctor.BiasAddLinearCombinationClamp,

  586.                     EpilogueFunctor.BiasAddLinearCombinationReluClamp,

  587.                     EpilogueFunctor.BiasAddLinearCombinationHSwishClamp,

  588.                 ]

  589.         elif (

  590.             tile.math_instruction.element_accumulator == DataType.f32

  591.             or tile.math_instruction.element_accumulator == DataType.f16

  592.         ) or (

  593.             tile.math_instruction.element_accumulator == DataType.s32

  594.             and out_dtype == DataType.f32

  595.         ):

  596.             bias_type = out_dtype

  597.             epilogues = [

  598.                 EpilogueFunctor.BiasAddLinearCombination,

  599.                 EpilogueFunctor.BiasAddLinearCombinationRelu,

  600.                 EpilogueFunctor.LinearCombination,

  601.             ]

  602.             if conv_type == ConvType.Convolution:

  603.                 epilogues.append(EpilogueFunctor.BiasAddLinearCombinationHSwish)

  604.         else:

  605.             assert False, "invalid path"

  606.         return bias_type, epilogues

  607. 

  608.     # rule for filter alignment

  609.     def get_flt_align(tile: TileDescription) -> int:

  610.         nonlocal flt_align

  611.         if (

  612.             tile.math_instruction.opcode_class == OpcodeClass.Simt

  613.             and tile.math_instruction.element_accumulator == DataType.s32

  614.         ):

  615.             thread_num = (

  616.                 tile.warp_count[0] * tile.warp_count[1] * tile.warp_count[2] * 32

  617.             )

  618.             flt_block = (

  619.                 tile.threadblock_shape[0]

  620.                 * tile.threadblock_shape[2]

  621.                 * DataTypeSize[tile.math_instruction.element_a]

  622.             )

  623.             load_per_thread = flt_block // thread_num

  624.             if load_per_thread >= 128:

  625.                 flt_align = 128

  626.             elif load_per_thread >= 64:

  627.                 flt_align = 64

  628.             else:

  629.                 assert load_per_thread >= 32

  630.                 flt_align = 32

  631.         return flt_align

  632. 

  633.     def get_dst_align(tile: TileDescription, out_layout: LayoutType) -> int:

  634.         nonlocal dst_align

  635.         if (

  636.             tile.math_instruction.opcode_class == OpcodeClass.TensorOp

  637.             and dst_layout == LayoutType.TensorNC4HW4

  638.         ):

  639.             dst_align = 32

  640.         return dst_align

  641. 

  642.     def filter_epilogue_with_conv_kind(

  643.         epilogue: EpilogueFunctor, conv_kind: ConvKind

  644.     ) -> bool:

  645.         if conv_kind == ConvKind.Fprop:

  646.             return epilogue == EpilogueFunctor.LinearCombination

  647.         elif conv_kind == ConvKind.Dgrad:

  648.             return (

  649.                 epilogue != EpilogueFunctor.BiasAddLinearCombinationClamp

  650.                 and epilogue != EpilogueFunctor.BiasAddLinearCombination

  651.             )

  652.         elif conv_kind == ConvKind.Wgrad:

  653.             return epilogue != EpilogueFunctor.LinearCombination

  654. 

  655.     # loop over all tile descriptions

  656.     for tile in tile_descriptions:

  657.         if filter_tile_with_layout(tile, dst_layout):

  658.             continue

  659. 

  660.         bias_type, epilogues = get_bias_type_and_epilogues(tile, dst_type)

  661. 

  662.         flt_align = flt_align if conv_kind == ConvKind.Wgrad else get_flt_align(tile)

  663. 

  664.         dst_align = get_dst_align(tile, dst_layout)

  665. 

  666.         for epilogue in epilogues:

  667.             if filter_epilogue_with_conv_kind(epilogue, conv_kind):

  668.                 continue

  669. 

  670.             if dst_type == DataType.f32:

  671.                 bias_type = DataType.f32

  672.             #

  673.             src = TensorDescription(

  674.                 tile.math_instruction.element_b,

  675.                 src_layout,

  676.                 int(src_align / DataTypeSize[tile.math_instruction.element_b]),

  677.             )

  678.             flt = TensorDescription(

  679.                 tile.math_instruction.element_a,

  680.                 flt_layout,

  681.                 int(flt_align / DataTypeSize[tile.math_instruction.element_a]),

  682.             )

  683.             bias = TensorDescription(

  684.                 bias_type, dst_layout, max(1, int(32 / DataTypeSize[bias_type]))

  685.             )

  686.             dst = TensorDescription(

  687.                 dst_type, dst_layout, int(dst_align / DataTypeSize[dst_type])

  688.             )

  689. 

  690.             new_operation = Conv2dOperation(

  691.                 conv_kind,

  692.                 conv_type,

  693.                 min_cc,

  694.                 tile,

  695.                 src,

  696.                 flt,

  697.                 bias,

  698.                 dst,

  699.                 element_epilogue,

  700.                 epilogue,

  701.                 swizzling_functor,

  702.                 SpecialOptimizeDesc.NoneSpecialOpt,

  703.                 implicit_gemm_mode,

  704.                 without_shared_load,

  705.                 required_cuda_ver_major,

  706.                 required_cuda_ver_minor,

  707.             )

  708.             operations.append(new_operation)

  709.             if use_special_optimization != SpecialOptimizeDesc.NoneSpecialOpt:

  710.                 new_operation = Conv2dOperation(

  711.                     conv_kind,

  712.                     conv_type,

  713.                     min_cc,

  714.                     tile,

  715.                     src,

  716.                     flt,

  717.                     bias,

  718.                     dst,

  719.                     element_epilogue,

  720.                     epilogue,

  721.                     swizzling_functor,

  722.                     use_special_optimization,

  723.                     implicit_gemm_mode,

  724.                     without_shared_load,

  725.                     required_cuda_ver_major,

  726.                     required_cuda_ver_minor,

  727.                 )

  728.                 operations.append(new_operation)

  729.     return operations

  730. 

  731. 

  732. ###################################################################################################

  733. #

  734. # Emitters functions for all targets

  735. #

  736. ###################################################################################################

  737. 

  738. 

  739. class EmitConv2dConfigurationLibrary:

  740.     def __init__(self, operation_path, configuration_name):

  741.         self.configuration_name = configuration_name

  742.         self.configuration_path = os.path.join(

  743.             operation_path, "%s.cu" % configuration_name

  744.         )

  745. 

  746.         self.instance_emitter = EmitConv2dInstance()

  747. 

  748.         self.instance_template = """

  749. ${operation_instance}

  750. 

  751. // Derived class

  752. struct ${operation_name} : 

  753.   public ${operation_name}_base { };

  754. 

  755. ///////////////////////////////////////////////////////////////////////////////////////////////////

  756. 

  757. """

  758.         self.header_template = """

  759. /*

  760.   Generated by conv2d_operation.py - Do not edit.

  761. */

  762. 

  763. ///////////////////////////////////////////////////////////////////////////////////////////////////

  764. 

  765. #include "cutlass/cutlass.h"

  766. #include "cutlass/library/library.h"

  767. #include "cutlass/library/manifest.h"

  768. 

  769. #include "library_internal.h"

  770. #include "conv2d_operation.h"

  771. 

  772. ///////////////////////////////////////////////////////////////////////////////////////////////////

  773. """

  774. 

  775.         self.configuration_header = """

  776. 

  777. namespace cutlass {

  778. namespace library {

  779. 

  780. // Initialize all instances

  781. void initialize_${configuration_name}(Manifest &manifest) {

  782. 

  783. """

  784. 

  785.         self.configuration_instance = """

  786.   using Operation_${operation_name} = cutlass::conv::device::ImplicitGemmConvolution<

  787.     ${operation_name}>;

  788. 

  789.   manifest.append(new cutlass::library::Conv2dOperation<

  790.     Operation_${operation_name}>(

  791.       "${operation_name}"));

  792. 

  793. """

  794. 

  795.         self.configuration_epilogue = """

  796. }

  797. """

  798.         self.epilogue_template = """

  799. 

  800. ///////////////////////////////////////////////////////////////////////////////////////////////////

  801. 

  802. } // namespace library

  803. } // namespace cutlass

  804. 

  805. ///////////////////////////////////////////////////////////////////////////////////////////////////

  806. 

  807. """

  808. 

  809.     #

  810.     def __enter__(self):

  811.         self.configuration_file = open(self.configuration_path, "w")

  812.         self.configuration_file.write(

  813.             SubstituteTemplate(

  814.                 self.header_template, {"configuration_name": self.configuration_name}

  815.             )

  816.         )

  817.         self.operations = []

  818.         return self

  819. 

  820.     #

  821.     def emit(self, operation):

  822.         self.operations.append(operation)

  823.         self.configuration_file.write(

  824.             SubstituteTemplate(

  825.                 self.instance_template,

  826.                 {

  827.                     "configuration_name": self.configuration_name,

  828.                     "operation_name": operation.procedural_name(),

  829.                     "operation_instance": self.instance_emitter.emit(operation),

  830.                 },

  831.             )

  832.         )

  833. 

  834.     #

  835.     def __exit__(self, exception_type, exception_value, traceback):

  836. 

  837.         self.configuration_file.write(

  838.             SubstituteTemplate(

  839.                 self.configuration_header,

  840.                 {"configuration_name": self.configuration_name},

  841.             )

  842.         )

  843. 

  844.         for operation in self.operations:

  845.             self.configuration_file.write(

  846.                 SubstituteTemplate(

  847.                     self.configuration_instance,

  848.                     {

  849.                         "configuration_name": self.configuration_name,

  850.                         "operation_name": operation.procedural_name(),

  851.                     },

  852.                 )

  853.             )

  854. 

  855.         self.configuration_file.write(self.configuration_epilogue)

  856.         self.configuration_file.write(self.epilogue_template)

  857.         self.configuration_file.close()

  858. 

  859. 

  860. ###################################################################################################

  861. ###################################################################################################

  862. 

  863. # Emitters for Conv Kernel Wrapper

  864. #

  865. ###################################################################################################

  866. 

  867. 

  868. class EmitConvSingleKernelWrapper:

  869.     def __init__(self, kernel_path, operation, short_path=False):

  870.         self.kernel_path = kernel_path

  871.         self.operation = operation

  872.         self.short_path = short_path

  873. 

  874.         if self.operation.conv_kind == ConvKind.Fprop:

  875.             self.instance_emitter = EmitConv2dInstance()

  876.             self.convolution_name = "ConvolutionOperation"

  877.         elif self.operation.conv_kind == ConvKind.Dgrad:

  878.             self.instance_emitter = EmitDeconvInstance()

  879.             self.convolution_name = "ConvolutionOperation"

  880.         else:

  881.             assert self.operation.conv_kind == ConvKind.Wgrad

  882.             self.instance_emitter = EmitConvolutionBackwardFilterInstance()

  883.             self.convolution_name = "ConvolutionBackwardFilterOperation"

  884. 

  885.         self.header_template = """

  886. #if __CUDACC_VER_MAJOR__ > ${required_cuda_ver_major} || (__CUDACC_VER_MAJOR__ == ${required_cuda_ver_major} && __CUDACC_VER_MINOR__ >= ${required_cuda_ver_minor})

  887. // ignore warning of cutlass

  888. #pragma GCC diagnostic push

  889. #pragma GCC diagnostic ignored "-Wunused-parameter"

  890. #pragma GCC diagnostic ignored "-Wstrict-aliasing"

  891. #pragma GCC diagnostic ignored "-Wuninitialized"

  892. #pragma GCC diagnostic ignored "-Wmaybe-uninitialized"

  893. 

  894. #include "cutlass/convolution/device/convolution.h"

  895. 

  896. #include "src/cuda/cutlass/manifest.h"

  897. #include "src/cuda/cutlass/convolution_operation.h"

  898. """

  899.         self.instance_template = """

  900. ${operation_instance}

  901. """

  902. 

  903.         self.manifest_template = """

  904. namespace cutlass {

  905. namespace library {

  906. 

  907. void initialize_${operation_name}(Manifest &manifest) {

  908.   manifest.append(new ${convolution_name}<Convolution>(

  909.     "${operation_name}"

  910.   ));

  911. }

  912. 

  913. }  // namespace library

  914. }  // namespace cutlass

  915. """

  916. 

  917.         self.epilogue_template = """

  918. #pragma GCC diagnostic pop

  919. #endif

  920. """

  921. 

  922.     #

  923.     def __enter__(self):

  924.         if self.short_path:

  925.             self.kernel_path = os.path.join(self.kernel_path, "%s.cu" % GlobalCnt.cnt)

  926.             GlobalCnt.cnt += 1

  927.         else:

  928.             self.kernel_path = os.path.join(

  929.                 self.kernel_path, "%s.cu" % self.operation.procedural_name()

  930.             )

  931.         self.kernel_file = open(self.kernel_path, "w")

  932.         self.kernel_file.write(

  933.             SubstituteTemplate(

  934.                 self.header_template,

  935.                 {

  936.                     "required_cuda_ver_major": str(

  937.                         self.operation.required_cuda_ver_major

  938.                     ),

  939.                     "required_cuda_ver_minor": str(

  940.                         self.operation.required_cuda_ver_minor

  941.                     ),

  942.                 },

  943.             )

  944.         )

  945.         return self

  946. 

  947.     #

  948.     def emit(self):

  949.         self.kernel_file.write(

  950.             SubstituteTemplate(

  951.                 self.instance_template,

  952.                 {"operation_instance": self.instance_emitter.emit(self.operation)},

  953.             )

  954.         )

  955. 

  956.         # emit manifest helper

  957.         manifest = SubstituteTemplate(

  958.             self.manifest_template,

  959.             {

  960.                 "operation_name": self.operation.procedural_name(),

  961.                 "convolution_name": self.convolution_name,

  962.             },

  963.         )

  964.         self.kernel_file.write(manifest)

  965. 

  966.     #

  967.     def __exit__(self, exception_type, exception_value, traceback):

  968.         self.kernel_file.write(self.epilogue_template)

  969.         self.kernel_file.close()

  970. 

  971. 

  972. ###################################################################################################

  973. ###################################################################################################