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/library.py

695 lines
22 kB
Raw Blame History 

  1. #

  2. # \file generator.py

  3. #

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

  5. #

  6. 

  7. import re

  8. 

  9. ###################################################################################################

  10. 

  11. import enum

  12. 

  13. # The following block implements enum.auto() for Python 3.5 variants that don't include it such

  14. # as the default 3.5.2 on Ubuntu 16.04.

  15. #

  16. # https://codereview.stackexchange.com/questions/177309/reimplementing-pythons-enum-auto-for-compatibility

  17. 

  18. try:

  19.     from enum import auto as enum_auto

  20. except ImportError:

  21.     __cutlass_library_auto_enum = 0

  22. 

  23.     def enum_auto() -> int:

  24.         global __cutlass_library_auto_enum

  25.         i = __cutlass_library_auto_enum

  26.         __cutlass_library_auto_enum += 1

  27.         return i

  28. 

  29. 

  30. ###################################################################################################

  31. 

  32. #

  33. class GeneratorTarget(enum.Enum):

  34.     Library = enum_auto()

  35. 

  36. 

  37. #

  38. GeneratorTargetNames = {GeneratorTarget.Library: "library"}

  39. #

  40. 

  41. ###################################################################################################

  42. 

  43. #

  44. class DataType(enum.Enum):

  45.     b1 = enum_auto()

  46.     u4 = enum_auto()

  47.     u8 = enum_auto()

  48.     u16 = enum_auto()

  49.     u32 = enum_auto()

  50.     u64 = enum_auto()

  51.     s4 = enum_auto()

  52.     s8 = enum_auto()

  53.     s16 = enum_auto()

  54.     s32 = enum_auto()

  55.     s64 = enum_auto()

  56.     f16 = enum_auto()

  57.     bf16 = enum_auto()

  58.     f32 = enum_auto()

  59.     tf32 = enum_auto()

  60.     f64 = enum_auto()

  61.     cf16 = enum_auto()

  62.     cbf16 = enum_auto()

  63.     cf32 = enum_auto()

  64.     ctf32 = enum_auto()

  65.     cf64 = enum_auto()

  66.     cs4 = enum_auto()

  67.     cs8 = enum_auto()

  68.     cs16 = enum_auto()

  69.     cs32 = enum_auto()

  70.     cs64 = enum_auto()

  71.     cu4 = enum_auto()

  72.     cu8 = enum_auto()

  73.     cu16 = enum_auto()

  74.     cu32 = enum_auto()

  75.     cu64 = enum_auto()

  76.     invalid = enum_auto()

  77. 

  78. 

  79. #

  80. ShortDataTypeNames = {

  81.     DataType.s32: "i",

  82.     DataType.f16: "h",

  83.     DataType.f32: "s",

  84.     DataType.f64: "d",

  85.     DataType.cf32: "c",

  86.     DataType.cf64: "z",

  87. }

  88. 

  89. #

  90. DataTypeNames = {

  91.     DataType.b1: "b1",

  92.     DataType.u4: "u4",

  93.     DataType.u8: "u8",

  94.     DataType.u16: "u16",

  95.     DataType.u32: "u32",

  96.     DataType.u64: "u64",

  97.     DataType.s4: "s4",

  98.     DataType.s8: "s8",

  99.     DataType.s16: "s16",

  100.     DataType.s32: "s32",

  101.     DataType.s64: "s64",

  102.     DataType.f16: "f16",

  103.     DataType.bf16: "bf16",

  104.     DataType.f32: "f32",

  105.     DataType.tf32: "tf32",

  106.     DataType.f64: "f64",

  107.     DataType.cf16: "cf16",

  108.     DataType.cbf16: "cbf16",

  109.     DataType.cf32: "cf32",

  110.     DataType.ctf32: "ctf32",

  111.     DataType.cf64: "cf64",

  112.     DataType.cu4: "cu4",

  113.     DataType.cu8: "cu8",

  114.     DataType.cu16: "cu16",

  115.     DataType.cu32: "cu32",

  116.     DataType.cu64: "cu64",

  117.     DataType.cs4: "cs4",

  118.     DataType.cs8: "cs8",

  119.     DataType.cs16: "cs16",

  120.     DataType.cs32: "cs32",

  121.     DataType.cs64: "cs64",

  122. }

  123. 

  124. DataTypeTag = {

  125.     DataType.b1: "cutlass::uint1b_t",

  126.     DataType.u4: "cutlass::uint4b_t",

  127.     DataType.u8: "uint8_t",

  128.     DataType.u16: "uint16_t",

  129.     DataType.u32: "uint32_t",

  130.     DataType.u64: "uint64_t",

  131.     DataType.s4: "cutlass::int4b_t",

  132.     DataType.s8: "int8_t",

  133.     DataType.s16: "int16_t",

  134.     DataType.s32: "int32_t",

  135.     DataType.s64: "int64_t",

  136.     DataType.f16: "cutlass::half_t",

  137.     DataType.bf16: "cutlass::bfloat16_t",

  138.     DataType.f32: "float",

  139.     DataType.tf32: "cutlass::tfloat32_t",

  140.     DataType.f64: "double",

  141.     DataType.cf16: "cutlass::complex<cutlass::half_t>",

  142.     DataType.cbf16: "cutlass::complex<cutlass::bfloat16_t>",

  143.     DataType.cf32: "cutlass::complex<float>",

  144.     DataType.ctf32: "cutlass::complex<cutlass::tfloat32_t>",

  145.     DataType.cf64: "cutlass::complex<double>",

  146.     DataType.cu4: "cutlass::complex<cutlass::uint4b_t>",

  147.     DataType.cu8: "cutlass::complex<cutlass::uint8_t>",

  148.     DataType.cu16: "cutlass::complex<cutlass::uint16_t>",

  149.     DataType.cu32: "cutlass::complex<cutlass::uint32_t>",

  150.     DataType.cu64: "cutlass::complex<cutlass::uint64_t>",

  151.     DataType.cs4: "cutlass::complex<cutlass::int4b_t>",

  152.     DataType.cs8: "cutlass::complex<cutlass::int8_t>",

  153.     DataType.cs16: "cutlass::complex<cutlass::int16_t>",

  154.     DataType.cs32: "cutlass::complex<cutlass::int32_t>",

  155.     DataType.cs64: "cutlass::complex<cutlass::int64_t>",

  156. }

  157. 

  158. DataTypeSize = {

  159.     DataType.b1: 1,

  160.     DataType.u4: 4,

  161.     DataType.u8: 4,

  162.     DataType.u16: 16,

  163.     DataType.u32: 32,

  164.     DataType.u64: 64,

  165.     DataType.s4: 4,

  166.     DataType.s8: 8,

  167.     DataType.s16: 16,

  168.     DataType.s32: 32,

  169.     DataType.s64: 64,

  170.     DataType.f16: 16,

  171.     DataType.bf16: 16,

  172.     DataType.f32: 32,

  173.     DataType.tf32: 32,

  174.     DataType.f64: 64,

  175.     DataType.cf16: 32,

  176.     DataType.cbf16: 32,

  177.     DataType.cf32: 64,

  178.     DataType.ctf32: 32,

  179.     DataType.cf64: 128,

  180.     DataType.cu4: 8,

  181.     DataType.cu8: 16,

  182.     DataType.cu16: 32,

  183.     DataType.cu32: 64,

  184.     DataType.cu64: 128,

  185.     DataType.cs4: 8,

  186.     DataType.cs8: 16,

  187.     DataType.cs16: 32,

  188.     DataType.cs32: 64,

  189.     DataType.cs64: 128,

  190. }

  191. 

  192. ###################################################################################################

  193. 

  194. #

  195. class ComplexTransform(enum.Enum):

  196.     none = enum_auto()

  197.     conj = enum_auto()

  198. 

  199. 

  200. #

  201. ComplexTransformTag = {

  202.     ComplexTransform.none: "cutlass::ComplexTransform::kNone",

  203.     ComplexTransform.conj: "cutlass::ComplexTransform::kConjugate",

  204. }

  205. 

  206. #

  207. RealComplexBijection = [

  208.     (DataType.f16, DataType.cf16),

  209.     (DataType.f32, DataType.cf32),

  210.     (DataType.f64, DataType.cf64),

  211. ]

  212. 

  213. #

  214. def is_complex(data_type):

  215.     for r, c in RealComplexBijection:

  216.         if data_type == c:

  217.             return True

  218.     return False

  219. 

  220. 

  221. #

  222. def get_complex_from_real(real_type):

  223.     for r, c in RealComplexBijection:

  224.         if real_type == r:

  225.             return c

  226.     return DataType.invalid

  227. 

  228. 

  229. #

  230. def get_real_from_complex(complex_type):

  231.     for r, c in RealComplexBijection:

  232.         if complex_type == c:

  233.             return r

  234.     return DataType.invalid

  235. 

  236. 

  237. #

  238. class ComplexMultiplyOp(enum.Enum):

  239.     multiply_add = enum_auto()

  240.     gaussian = enum_auto()

  241. 

  242. 

  243. ###################################################################################################

  244. 

  245. #

  246. class MathOperation(enum.Enum):

  247.     multiply_add = enum_auto()

  248.     multiply_add_saturate = enum_auto()

  249.     xor_popc = enum_auto()

  250.     multiply_add_fast_bf16 = enum_auto()

  251.     multiply_add_fast_f16 = enum_auto()

  252.     multiply_add_complex = enum_auto()

  253.     multiply_add_complex_gaussian = enum_auto()

  254. 

  255. 

  256. #

  257. MathOperationTag = {

  258.     MathOperation.multiply_add: "cutlass::arch::OpMultiplyAdd",

  259.     MathOperation.multiply_add_saturate: "cutlass::arch::OpMultiplyAddSaturate",

  260.     MathOperation.xor_popc: "cutlass::arch::OpXorPopc",

  261.     MathOperation.multiply_add_fast_bf16: "cutlass::arch::OpMultiplyAddFastBF16",

  262.     MathOperation.multiply_add_fast_f16: "cutlass::arch::OpMultiplyAddFastF16",

  263.     MathOperation.multiply_add_complex: "cutlass::arch::OpMultiplyAddComplex",

  264.     MathOperation.multiply_add_complex_gaussian: "cutlass::arch::OpMultiplyAddGaussianComplex",

  265. }

  266. 

  267. ###################################################################################################

  268. 

  269. #

  270. class LayoutType(enum.Enum):

  271.     ColumnMajor = enum_auto()

  272.     RowMajor = enum_auto()

  273.     ColumnMajorInterleaved2 = enum_auto()

  274.     RowMajorInterleaved2 = enum_auto()

  275.     ColumnMajorInterleaved32 = enum_auto()

  276.     RowMajorInterleaved32 = enum_auto()

  277.     ColumnMajorInterleaved64 = enum_auto()

  278.     RowMajorInterleaved64 = enum_auto()

  279.     TensorNHWC = enum_auto()

  280.     TensorNDHWC = enum_auto()

  281.     TensorNCHW = enum_auto()

  282.     TensorNGHWC = enum_auto()

  283.     TensorNC4HW4 = enum_auto()

  284.     TensorC4RSK4 = enum_auto()

  285.     TensorNC8HW8 = enum_auto()

  286.     TensorNC16HW16 = enum_auto()

  287.     TensorNC32HW32 = enum_auto()

  288.     TensorNC64HW64 = enum_auto()

  289.     TensorC32RSK32 = enum_auto()

  290.     TensorC64RSK64 = enum_auto()

  291.     TensorK4RSC4 = enum_auto()

  292.     TensorCK4RS4 = enum_auto()

  293.     TensorCK8RS8 = enum_auto()

  294.     TensorCK16RS16 = enum_auto()

  295. 

  296. 

  297. #

  298. LayoutTag = {

  299.     LayoutType.ColumnMajor: "cutlass::layout::ColumnMajor",

  300.     LayoutType.RowMajor: "cutlass::layout::RowMajor",

  301.     LayoutType.ColumnMajorInterleaved2: "cutlass::layout::ColumnMajorInterleaved<2>",

  302.     LayoutType.RowMajorInterleaved2: "cutlass::layout::RowMajorInterleaved<2>",

  303.     LayoutType.ColumnMajorInterleaved32: "cutlass::layout::ColumnMajorInterleaved<32>",

  304.     LayoutType.RowMajorInterleaved32: "cutlass::layout::RowMajorInterleaved<32>",

  305.     LayoutType.ColumnMajorInterleaved64: "cutlass::layout::ColumnMajorInterleaved<64>",

  306.     LayoutType.RowMajorInterleaved64: "cutlass::layout::RowMajorInterleaved<64>",

  307.     LayoutType.TensorNHWC: "cutlass::layout::TensorNHWC",

  308.     LayoutType.TensorNDHWC: "cutlass::layout::TensorNDHWC",

  309.     LayoutType.TensorNCHW: "cutlass::layout::TensorNCHW",

  310.     LayoutType.TensorNGHWC: "cutlass::layout::TensorNGHWC",

  311.     LayoutType.TensorNC4HW4: "cutlass::layout::TensorNCxHWx<4>",

  312.     LayoutType.TensorC4RSK4: "cutlass::layout::TensorCxRSKx<4>",

  313.     LayoutType.TensorNC8HW8: "cutlass::layout::TensorNCxHWx<8>",

  314.     LayoutType.TensorNC16HW16: "cutlass::layout::TensorNCxHWx<16>",

  315.     LayoutType.TensorNC32HW32: "cutlass::layout::TensorNCxHWx<32>",

  316.     LayoutType.TensorC32RSK32: "cutlass::layout::TensorCxRSKx<32>",

  317.     LayoutType.TensorNC64HW64: "cutlass::layout::TensorNCxHWx<64>",

  318.     LayoutType.TensorC64RSK64: "cutlass::layout::TensorCxRSKx<64>",

  319.     LayoutType.TensorK4RSC4: "cutlass::layout::TensorKxRSCx<4>",

  320.     LayoutType.TensorCK4RS4: "cutlass::layout::TensorCKxRSx<4>",

  321.     LayoutType.TensorCK8RS8: "cutlass::layout::TensorCKxRSx<8>",

  322.     LayoutType.TensorCK16RS16: "cutlass::layout::TensorCKxRSx<16>",

  323. }

  324. 

  325. #

  326. TransposedLayout = {

  327.     LayoutType.ColumnMajor: LayoutType.RowMajor,

  328.     LayoutType.RowMajor: LayoutType.ColumnMajor,

  329.     LayoutType.ColumnMajorInterleaved2: LayoutType.RowMajorInterleaved2,

  330.     LayoutType.RowMajorInterleaved2: LayoutType.ColumnMajorInterleaved2,

  331.     LayoutType.ColumnMajorInterleaved32: LayoutType.RowMajorInterleaved32,

  332.     LayoutType.RowMajorInterleaved32: LayoutType.ColumnMajorInterleaved32,

  333.     LayoutType.ColumnMajorInterleaved64: LayoutType.RowMajorInterleaved64,

  334.     LayoutType.RowMajorInterleaved64: LayoutType.ColumnMajorInterleaved64,

  335.     LayoutType.TensorNHWC: LayoutType.TensorNHWC,

  336. }

  337. 

  338. #

  339. ShortLayoutTypeNames = {

  340.     LayoutType.ColumnMajor: "n",

  341.     LayoutType.ColumnMajorInterleaved32: "n2",

  342.     LayoutType.ColumnMajorInterleaved32: "n32",

  343.     LayoutType.ColumnMajorInterleaved64: "n64",

  344.     LayoutType.RowMajor: "t",

  345.     LayoutType.RowMajorInterleaved2: "t2",

  346.     LayoutType.RowMajorInterleaved32: "t32",

  347.     LayoutType.RowMajorInterleaved64: "t64",

  348.     LayoutType.TensorNHWC: "nhwc",

  349.     LayoutType.TensorNDHWC: "ndhwc",

  350.     LayoutType.TensorNCHW: "nchw",

  351.     LayoutType.TensorNGHWC: "nghwc",

  352.     LayoutType.TensorNC4HW4: "nc4hw4",

  353.     LayoutType.TensorC4RSK4: "c4rsk4",

  354.     LayoutType.TensorNC8HW8: "nc8hw8",

  355.     LayoutType.TensorNC16HW16: "nc16hw16",

  356.     LayoutType.TensorNC32HW32: "nc32hw32",

  357.     LayoutType.TensorNC64HW64: "nc64hw64",

  358.     LayoutType.TensorC32RSK32: "c32rsk32",

  359.     LayoutType.TensorC64RSK64: "c64rsk64",

  360.     LayoutType.TensorK4RSC4: "k4rsc4",

  361.     LayoutType.TensorCK4RS4: "ck4rs4",

  362.     LayoutType.TensorCK8RS8: "ck8rs8",

  363.     LayoutType.TensorCK16RS16: "ck16rs16",

  364. }

  365. 

  366. #

  367. ShortComplexLayoutNames = {

  368.     (LayoutType.ColumnMajor, ComplexTransform.none): "n",

  369.     (LayoutType.ColumnMajor, ComplexTransform.conj): "c",

  370.     (LayoutType.RowMajor, ComplexTransform.none): "t",

  371.     (LayoutType.RowMajor, ComplexTransform.conj): "h",

  372. }

  373. 

  374. ###################################################################################################

  375. #

  376. class OpcodeClass(enum.Enum):

  377.     Simt = enum_auto()

  378.     TensorOp = enum_auto()

  379.     WmmaTensorOp = enum_auto()

  380. 

  381. 

  382. OpcodeClassNames = {

  383.     OpcodeClass.Simt: "simt",

  384.     OpcodeClass.TensorOp: "tensorop",

  385.     OpcodeClass.WmmaTensorOp: "wmma_tensorop",

  386. }

  387. 

  388. OpcodeClassTag = {

  389.     OpcodeClass.Simt: "cutlass::arch::OpClassSimt",

  390.     OpcodeClass.TensorOp: "cutlass::arch::OpClassTensorOp",

  391.     OpcodeClass.WmmaTensorOp: "cutlass::arch::OpClassWmmaTensorOp",

  392. }

  393. 

  394. ###################################################################################################

  395. 

  396. #

  397. class OperationKind(enum.Enum):

  398.     Gemm = enum_auto()

  399.     Conv2d = enum_auto()

  400. 

  401. 

  402. #

  403. OperationKindNames = {OperationKind.Gemm: "gemm", OperationKind.Conv2d: "conv2d"}

  404. 

  405. #

  406. class Target(enum.Enum):

  407.     library = enum_auto()

  408. 

  409. 

  410. ArchitectureNames = {

  411.     50: "maxwell",

  412.     60: "pascal",

  413.     61: "pascal",

  414.     70: "volta",

  415.     75: "turing",

  416.     80: "ampere",

  417. }

  418. 

  419. ###################################################################################################

  420. 

  421. #

  422. def SubstituteTemplate(template, values):

  423.     text = template

  424.     changed = True

  425.     while changed:

  426.         changed = False

  427.         for key, value in values.items():

  428.             regex = "\\$\\{%s\\}" % key

  429.             newtext = re.sub(regex, value, text)

  430.             if newtext != text:

  431.                 changed = True

  432.             text = newtext

  433.     return text

  434. 

  435. 

  436. ###################################################################################################

  437. 

  438. #

  439. class GemmKind(enum.Enum):

  440.     Gemm = enum_auto()

  441.     Sparse = enum_auto()

  442.     Universal = enum_auto()

  443.     PlanarComplex = enum_auto()

  444.     PlanarComplexArray = enum_auto()

  445.     SplitKParallel = enum_auto()

  446.     GemvBatchedStrided = enum_auto()

  447. 

  448. 

  449. #

  450. GemmKindNames = {

  451.     GemmKind.Gemm: "gemm",

  452.     GemmKind.Sparse: "spgemm",

  453.     GemmKind.Universal: "gemm",

  454.     GemmKind.PlanarComplex: "gemm_planar_complex",

  455.     GemmKind.PlanarComplexArray: "gemm_planar_complex_array",

  456.     GemmKind.SplitKParallel: "gemm_split_k_parallel",

  457.     GemmKind.GemvBatchedStrided: "gemv_batched_strided",

  458. }

  459. 

  460. #

  461. class EpilogueFunctor(enum.Enum):

  462.     LinearCombination = enum_auto()

  463.     LinearCombinationClamp = enum_auto()

  464.     BiasAddLinearCombination = enum_auto()

  465.     BiasAddLinearCombinationRelu = enum_auto()

  466.     BiasAddLinearCombinationHSwish = enum_auto()

  467.     BiasAddLinearCombinationClamp = enum_auto()

  468.     BiasAddLinearCombinationReluClamp = enum_auto()

  469.     BiasAddLinearCombinationHSwishClamp = enum_auto()

  470. 

  471. 

  472. #

  473. EpilogueFunctorTag = {

  474.     EpilogueFunctor.LinearCombination: "cutlass::epilogue::thread::LinearCombination",

  475.     EpilogueFunctor.LinearCombinationClamp: "cutlass::epilogue::thread::LinearCombinationClamp",

  476.     EpilogueFunctor.BiasAddLinearCombination: "cutlass::epilogue::thread::BiasAddLinearCombination",

  477.     EpilogueFunctor.BiasAddLinearCombinationRelu: "cutlass::epilogue::thread::BiasAddLinearCombinationRelu",

  478.     EpilogueFunctor.BiasAddLinearCombinationHSwish: "cutlass::epilogue::thread::BiasAddLinearCombinationHSwish",

  479.     EpilogueFunctor.BiasAddLinearCombinationClamp: "cutlass::epilogue::thread::BiasAddLinearCombinationClamp",

  480.     EpilogueFunctor.BiasAddLinearCombinationReluClamp: "cutlass::epilogue::thread::BiasAddLinearCombinationReluClamp",

  481.     EpilogueFunctor.BiasAddLinearCombinationHSwishClamp: "cutlass::epilogue::thread::BiasAddLinearCombinationHSwishClamp",

  482. }

  483. 

  484. #

  485. ShortEpilogueNames = {

  486.     EpilogueFunctor.LinearCombination: "id",

  487.     EpilogueFunctor.BiasAddLinearCombinationHSwishClamp: "hswish",

  488.     EpilogueFunctor.BiasAddLinearCombinationReluClamp: "relu",

  489.     EpilogueFunctor.BiasAddLinearCombinationClamp: "id",

  490.     EpilogueFunctor.BiasAddLinearCombinationHSwish: "hswish",

  491.     EpilogueFunctor.BiasAddLinearCombinationRelu: "relu",

  492.     EpilogueFunctor.BiasAddLinearCombination: "id",

  493. }

  494. 

  495. 

  496. #

  497. class SwizzlingFunctor(enum.Enum):

  498.     Identity1 = enum_auto()

  499.     Identity2 = enum_auto()

  500.     Identity4 = enum_auto()

  501.     Identity8 = enum_auto()

  502.     ConvFpropNCxHWx = enum_auto()

  503.     ConvFpropTrans = enum_auto()

  504.     ConvDgradNCxHWx = enum_auto()

  505.     ConvDgradTrans = enum_auto()

  506.     DepthwiseConvolutionFprop = enum_auto()

  507.     DepthwiseConvolutionDgrad = enum_auto()

  508.     DepthwiseConvolutionWgrad = enum_auto()

  509. 

  510. 

  511. #

  512. SwizzlingFunctorTag = {

  513.     SwizzlingFunctor.Identity1: "cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<1>",

  514.     SwizzlingFunctor.Identity2: "cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<2>",

  515.     SwizzlingFunctor.Identity4: "cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<4>",

  516.     SwizzlingFunctor.Identity8: "cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>",

  517.     SwizzlingFunctor.ConvFpropNCxHWx: "cutlass::conv::threadblock::ConvolutionFpropNCxHWxThreadblockSwizzle",

  518.     SwizzlingFunctor.ConvFpropTrans: "cutlass::conv::threadblock::ConvolutionFpropTransThreadblockSwizzle",

  519.     SwizzlingFunctor.ConvDgradNCxHWx: "cutlass::conv::threadblock::ConvolutionDgradNCxHWxThreadblockSwizzle",

  520.     SwizzlingFunctor.ConvDgradTrans: "cutlass::conv::threadblock::ConvolutionDgradTransThreadblockSwizzle",

  521.     SwizzlingFunctor.DepthwiseConvolutionFprop: "cutlass::conv::threadblock::DepthwiseConvolutionFpropThreadblockSwizzle",

  522.     SwizzlingFunctor.DepthwiseConvolutionDgrad: "cutlass::conv::threadblock::DepthwiseConvolutionDgradThreadblockSwizzle",

  523.     SwizzlingFunctor.DepthwiseConvolutionWgrad: "cutlass::conv::threadblock::DepthwiseConvolutionWgradThreadblockSwizzle",

  524. }

  525. 

  526. ###################################################################################################

  527. 

  528. 

  529. class ConvType(enum.Enum):

  530.     Convolution = enum_auto()

  531.     BatchConvolution = enum_auto()

  532.     Local = enum_auto()

  533.     LocalShare = enum_auto()

  534.     DepthwiseConvolution = enum_auto()

  535. 

  536. 

  537. ConvTypeTag = {

  538.     ConvType.Convolution: "cutlass::conv::ConvType::kConvolution",

  539.     ConvType.BatchConvolution: "cutlass::conv::ConvType::kBatchConvolution",

  540.     ConvType.Local: "cutlass::conv::ConvType::kLocal",

  541.     ConvType.LocalShare: "cutlass::conv::ConvType::kLocalShare",

  542.     ConvType.DepthwiseConvolution: "cutlass::conv::ConvType::kDepthwiseConvolution",

  543. }

  544. 

  545. #

  546. class ConvKind(enum.Enum):

  547.     Fprop = enum_auto()

  548.     Dgrad = enum_auto()

  549.     Wgrad = enum_auto()

  550. 

  551. 

  552. #

  553. ConvKindTag = {

  554.     ConvKind.Fprop: "cutlass::conv::Operator::kFprop",

  555.     ConvKind.Dgrad: "cutlass::conv::Operator::kDgrad",

  556.     ConvKind.Wgrad: "cutlass::conv::Operator::kWgrad",

  557. }

  558. 

  559. ConvKindNames = {

  560.     ConvKind.Fprop: "fprop",

  561.     ConvKind.Dgrad: "dgrad",

  562.     ConvKind.Wgrad: "wgrad",

  563. }

  564. 

  565. #

  566. class IteratorAlgorithm(enum.Enum):

  567.     Analytic = enum_auto()

  568.     Optimized = enum_auto()

  569. 

  570. 

  571. #

  572. IteratorAlgorithmTag = {

  573.     IteratorAlgorithm.Analytic: "cutlass::conv::IteratorAlgorithm::kAnalytic",

  574.     IteratorAlgorithm.Optimized: "cutlass::conv::IteratorAlgorithm::kOptimized",

  575. }

  576. 

  577. IteratorAlgorithmNames = {

  578.     IteratorAlgorithm.Analytic: "analytic",

  579.     IteratorAlgorithm.Optimized: "optimized",

  580. }

  581. 

  582. #

  583. class StrideSupport(enum.Enum):

  584.     Strided = enum_auto()

  585.     Unity = enum_auto()

  586. 

  587. 

  588. #

  589. StrideSupportTag = {

  590.     StrideSupport.Strided: "cutlass::conv::StrideSupport::kStrided",

  591.     StrideSupport.Unity: "cutlass::conv::StrideSupport::kUnity",

  592. }

  593. 

  594. StrideSupportNames = {StrideSupport.Strided: "", StrideSupport.Unity: "unity_stride"}

  595. 

  596. 

  597. class SpecialOptimizeDesc(enum.Enum):

  598.     NoneSpecialOpt = enum_auto()

  599.     ConvFilterUnity = enum_auto()

  600.     DeconvDoubleUpsampling = enum_auto()

  601. 

  602. 

  603. SpecialOptimizeDescNames = {

  604.     SpecialOptimizeDesc.NoneSpecialOpt: "none",

  605.     SpecialOptimizeDesc.ConvFilterUnity: "conv_filter_unity",

  606.     SpecialOptimizeDesc.DeconvDoubleUpsampling: "deconv_double_upsampling",

  607. }

  608. 

  609. SpecialOptimizeDescTag = {

  610.     SpecialOptimizeDesc.NoneSpecialOpt: "cutlass::conv::SpecialOptimizeDesc::NONE",

  611.     SpecialOptimizeDesc.ConvFilterUnity: "cutlass::conv::SpecialOptimizeDesc::CONV_FILTER_UNITY",

  612.     SpecialOptimizeDesc.DeconvDoubleUpsampling: "cutlass::conv::SpecialOptimizeDesc::DECONV_DOUBLE_UPSAMPLING",

  613. }

  614. 

  615. 

  616. class ImplicitGemmMode(enum.Enum):

  617.     GemmNT = enum_auto()

  618.     GemmTN = enum_auto()

  619. 

  620. 

  621. ImplicitGemmModeNames = {

  622.     ImplicitGemmMode.GemmNT: "gemm_nt",

  623.     ImplicitGemmMode.GemmTN: "gemm_tn",

  624. }

  625. 

  626. ImplicitGemmModeTag = {

  627.     ImplicitGemmMode.GemmNT: "cutlass::conv::ImplicitGemmMode::GEMM_NT",

  628.     ImplicitGemmMode.GemmTN: "cutlass::conv::ImplicitGemmMode::GEMM_TN",

  629. }

  630. 

  631. ###################################################################################################

  632. 

  633. #

  634. class MathInstruction:

  635.     def __init__(

  636.         self,

  637.         instruction_shape,

  638.         element_a,

  639.         element_b,

  640.         element_accumulator,

  641.         opcode_class,

  642.         math_operation=MathOperation.multiply_add,

  643.     ):

  644.         self.instruction_shape = instruction_shape

  645.         self.element_a = element_a

  646.         self.element_b = element_b

  647.         self.element_accumulator = element_accumulator

  648.         self.opcode_class = opcode_class

  649.         self.math_operation = math_operation

  650. 

  651. 

  652. #

  653. class TileDescription:

  654.     def __init__(

  655.         self,

  656.         threadblock_shape,

  657.         stages,

  658.         warp_count,

  659.         math_instruction,

  660.         min_compute,

  661.         max_compute,

  662.     ):

  663.         self.threadblock_shape = threadblock_shape

  664.         self.stages = stages

  665.         self.warp_count = warp_count

  666.         self.math_instruction = math_instruction

  667.         self.minimum_compute_capability = min_compute

  668.         self.maximum_compute_capability = max_compute

  669. 

  670.     def procedural_name(self):

  671.         return "%dx%d_%dx%d" % (

  672.             self.threadblock_shape[0],

  673.             self.threadblock_shape[1],

  674.             self.threadblock_shape[2],

  675.             self.stages,

  676.         )

  677. 

  678. 

  679. #

  680. class TensorDescription:

  681.     def __init__(

  682.         self, element, layout, alignment=1, complex_transform=ComplexTransform.none

  683.     ):

  684.         self.element = element

  685.         self.layout = layout

  686.         self.alignment = alignment

  687.         self.complex_transform = complex_transform

  688. 

  689. 

  690. ###################################################################################################

  691. 

  692. 

  693. class GlobalCnt:

  694.     cnt = 0