Huawei_Technology
/
MindSpore-AKG
Not watched
1
0
Fork 0
Code
Releases 1 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
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.
Branch: master
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'master'
${ noResults }
MindSpore-AKG/python/akg/ops/nn/batchmatmul.py

1187 lines
70 kB
Raw Permalink Blame History 

  1. #!/usr/bin/env python3

  2. # coding: utf-8

  3. # Copyright 2019 Huawei Technologies Co., Ltd

  4. #

  5. # Licensed under the Apache License, Version 2.0 (the "License");

  6. # you may not use this file except in compliance with the License.

  7. # You may obtain a copy of the License at

  8. #

  9. # http://www.apache.org/licenses/LICENSE-2.0

  10. #

  11. # Unless required by applicable law or agreed to in writing, software

  12. # distributed under the License is distributed on an "AS IS" BASIS,

  13. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

  14. # See the License for the specific language governing permissions and

  15. # limitations under the License.

  16. 

  17. """operator dsl function: batchmatmul"""

  18. from functools import reduce

  19. import akg.topi

  20. import akg.tvm

  21. from akg.tvm.hybrid import script

  22. from akg.ops.math import cast

  23. from akg.utils import custom_tiling as ct_util

  24. from akg.utils import validation_check as vc_util

  25. from akg.utils.format_transform import get_shape, get_bytes

  26. from akg.utils.math import greatest_common_divisor, least_common_multiple

  27. from akg.utils.kernel_exec import product_is_mini

  28. from akg.utils import dynamic_shape as ds

  29. 

  30. batchmatmul_set_dim_map = {

  31.     # 2D

  32.     str((256, 1024, 4096, "float32", False, True)): ((16, 16), (16, 16), (16, 16)),

  33.     str((160, 1024, 1024, "float32", False, False)): ((1, 1), (16, 16), (1024, 1024)),

  34.     str((8192, 1024, 4096, "float32", False, False)): ((1, 1), (1024, 1024), (16, 16)),

  35.     str((1024, 1024, 8192, "float32", True, False)): ((8, 8), (8, 8), (512, 512)),

  36.     str((1024, 1024, 2, "float32", False, False)): ((64, 64), (64, 64), (2, 2)),

  37.     str((1024, 1024, 4096, "float32", False, False)): ((1, 1), (16, 16), (512, 512)),

  38.     str((2, 1024, 8192, "float32", True, False)): ((2, 2), (16, 16), (512, 512)),

  39.     str((30522, 1024, 1280, "float32", True, False)): ((3, 3), (64, 64), (128, 128)),

  40.     str((1024, 4096, 8192, "float32", True, False)): ((32, 32), (32, 32), (32, 32)),

  41.     str((2, 1024, 64, "float32", True, False)): ((4, 4), (64, 64), (64, 64)),

  42.     str((160, 30522, 1024, "float32", False, True)): ((1, 1), (3, 3), (512, 512)),

  43.     str((1024, 1024, 64, "float32", True, False)): ((16, 16), (16, 16), (64, 64)),

  44.     str((4096, 1024, 8192, "float32", True, False)): ((16, 16), (64, 64), (16, 16)),

  45.     str((1280, 1024, 30522, "float32", False, False)): ((1, 1), (512, 512), (3, 3)),

  46.     str((8192, 1024, 4096, "float32", False, True)): ((1, 1), (16, 16), (1024, 1024)),

  47.     str((1280, 30522, 1024, "float32", False, True)): ((1, 1), (3, 3), (512, 512)),

  48.     str((8192, 4096, 1024, "float32", False, False)): ((1, 1), (64, 64), (256, 256)),

  49.     str((768, 768, 8192, "float16", True, False)): ((16, 16), (16, 16), (64, 64)),

  50.     str((3072, 768, 8192, "float16", True, False)): ((16, 16), (16, 16), (64, 64)),

  51.     str((2, 768, 64, "float32", True, False)): ((2, 2), (64, 64), (64, 64)),

  52.     str((768, 1024, 8192, "float16", False, False)): ((16, 1), (16, 1), (64, 1)),

  53.     str((33, 64, 16384, "float32", True, False)): ((1, 1), (16, 16), (128, 128)),

  54.     str((768, 768, 64, "float32", True, False)): ((16, 16), (16, 16), (16, 16)),

  55.     str((8192, 768, 21128, "float32", False, False)): ((4, 4), (128, 128), (4, 4)),

  56.     str((2, 768, 8192, "float32", True, False)): ((2, 2), (16, 16), (512, 512)),

  57.     str((8192, 768, 768, "float16", False, True)): ((1, 1), (16, 16), (768, 768)),

  58.     str((21128, 768, 8192, "float32", True, False)): ((4, 4), (128, 128), (64, 64)),

  59.     str((768, 1024, 768, "float32", True, False)): ((2, 2), (128, 128), (128, 128)),

  60.     str((16384, 16384, 33, "float32", True, False)): ((16, 16), (64, 64), (33, 33)),

  61.     str((21128, 768, 8192, "float32", False, False)): ((4, 4), (128, 128), (64, 64)),

  62.     str((1280, 1280, 1024, "float32", False, True)): ((4, 4), (32, 32), (128, 128)),

  63.     str((1280, 768, 21128, "float32", False, False)): ((1, 1), (768, 768), (8, 8)),

  64.     str((8192, 768, 768, "float32", False, False)): ((1, 1), (8, 8), (768, 768)),

  65.     str((20, 768, 32000, "float32", False, False)): ((20, 20), (48, 48), (32, 32)),

  66.     str((21128, 768, 1280, "float32", True, False)): ((2, 2), (32, 32), (32, 32)),

  67.     str((768, 3072, 1892, "float32", True, False)): ((16, 16), (16, 16), (16, 16)),

  68.     str((33, 64, 16384, "float32", False, True)): ((2, 2), (32, 32), (32, 32)),

  69.     str((8192, 3072, 768, "float32", False, True)): ((16, 16), (16, 16), (16, 16)),

  70.     str((2, 8192, 768, "float32", True, False)): ((16, 16), (16, 16), (16, 16)),

  71.     str((8192, 768, 3072, "float32", False, False)): ((32, 32), (32, 32), (32, 32)),

  72.     str((8192, 768, 3072, "float16", False, True)): ((1, 1), (16, 1), (768, 1)),

  73.     str((8192, 3072, 768, "float16", False, True)): ((1, 1), (16, 1), (768, 1)),

  74.     str((8192, 768, 768, "float16", False, False)): ((1, 1), (16, 1), (768, 1)),

  75.     str((768, 3072, 768, "float16", False, False)): ((1, 1), (16, 1), (768, 1)),

  76.     str((8192, 3072, 768, "float16", False, False)): ((1, 1), (16, 1), (768, 1)),

  77.     str((8192, 768, 3072, "float16", False, False)): ((1, 1), (16, 1), (768, 1)),

  78.     str((21128, 768, 5120, "float32", True, False)): ((4, 4), (128, 128), (64, 64)),

  79.     str((21128, 768, 2560, "float32", True, False)): ((4, 4), (128, 128), (64, 64)),

  80.     str((1024, 2, 4, "float32", True, False)): ((512, 1), (2, 1), (4, 1)),

  81.     str((21128, 1024, 21128, "float32", False, False)): ((4, 4), (512, 512), (4, 4)),

  82.     str((320, 768, 21128, "float32", False, False)): ((40, 40), (128, 128), (4, 4)),

  83.     str((5120, 1024, 21128, "float32", False, False)): ((64, 64), (128, 128), (4, 4)),

  84.     str((16384, 4096, 1024, "float32", False, False)): ((4, 4), (64, 64), (64, 64)),

  85.     str((1024, 4096, 16384, "float32", True, False)): ((64, 64), (64, 64), (4, 4)),

  86.     str((768, 3072, 8192, "float16", True, False)): ((1, 1), (16, 1), (512, 1)),

  87.     str((2048, 768, 3072, "float32", False, False)): ((8, 1), (16, 1), (128, 1)),

  88.     str((1024, 2, 128, "float32", True, False)): ((32, 1), (2, 1), (128, 1)),

  89.     str((1024, 2, 16, "float32", True, False)): ((256, 1), (2, 1), (16, 1)),

  90.     str((1024, 2, 32, "float32", True, False)): ((128, 1), (2, 1), (32, 1)),

  91.     str((1024, 2, 64, "float32", True, False)): ((64, 1), (2, 1), (64, 1)),

  92.     str((1024, 2, 8, "float32", True, False)): ((512, 1), (2, 1), (8, 1)),

  93.     str((768, 2, 128, "float32", True, False)): ((32, 1), (2, 1), (128, 1)),

  94.     str((768, 2, 16, "float32", True, False)): ((256, 1), (2, 1), (16, 1)),

  95.     str((768, 2, 32, "float32", True, False)): ((128, 1), (2, 1), (32, 1)),

  96.     str((768, 2, 64, "float32", True, False)): ((64, 1), (2, 1), (64, 1)),

  97.     str((768, 3072, 2048, "float32", True, False)): ((16, 1), (16, 1), (64, 1)),

  98.     str((3072, 768, 2048, "float32", True, False)): ((16, 1), (16, 1), (64, 1)),

  99.     str((65536, 1024, 4096, "float32", False, False)): ((1, 1), (1024, 1), (16, 1)),

  100.     str((10240, 1024, 21128, "float32", False, False)): ((1, 1), (1024, 1), (16, 1)),

  101.     str((10240, 21128, 1024, "float32", False, False)): ((1, 1), (1024, 1), (16, 1)),

  102.     str((21128, 768, 21128, "float32", True, False)): ((1, 1), (768, 1), (16, 1)),

  103.     str((2560, 21128, 1024, "float32", False, False)): ((1, 1), (1024, 1), (16, 1)),

  104.     str((5120, 21128, 1024, "float32", False, False)): ((1, 1), (1024, 1), (16, 1)),

  105.     str((32768, 4096, 1024, "float32", False, False)): ((1, 1), (1024, 1), (16, 1)),

  106.     str((20480, 1024, 21128, "float32", False, False)): ((1, 1), (1024, 1), (16, 1)),

  107.     str((128, 1024, 4096, "float32", False, True)): ((1, 1), (16, 1), (16, 1)),

  108.     str((128, 1024, 4096, "float16", False, True)): ((1, 1), (16, 1), (32, 1)),

  109.     str((128, 4096, 1024, "float16", False, True)): ((1, 1), (32, 1), (32, 1)),

  110.     str((20480, 21128, 1024, "float32", False, False)): ((1, 1), (1024, 1), (16, 1)),

  111.     str((65536, 4096, 1024, "float32", False, False)): ((1, 1), (1024, 1), (16, 1)),

  112.     str((512, 1024, 4096, "float32", False, True)): ((1, 1), (8, 1), (1024, 1)),

  113.     str((256, 1024, 4096, "float32", False, True)): ((1, 1), (8, 1), (1024, 1)),

  114.     str((1024, 1024, 4096, "float32", False, True)): ((1, 1), (8, 1), (1024, 1)),

  115.     str((2048, 1024, 4096, "float32", False, True)): ((1, 1), (8, 1), (1024, 1)),

  116.     str((2, 2, 1024, "float32", False, True)): ((1, 1), (2, 1), (1024, 1)),

  117.     str((65536, 768, 3072, "float32", False, True)): ((1, 1), (192, 1), (96, 1)),

  118.     str((16384, 1024, 4096, "float32", False, True)): ((1, 1), (128, 1), (128, 1)),

  119.     str((1024, 4096, 32768, "float32", True, False)): ((8, 1), (1024, 1), (4, 1)),

  120.     str((4, 2, 1024, "float32", False, True)): ((1, 1), (2, 1), (1024, 1)),

  121.     str((2, 2, 768, "float32", False, True)): ((1, 1), (2, 1), (768, 1)),

  122.     str((8, 2, 768, "float32", False, True)): ((1, 1), (2, 1), (768, 1)),

  123.     str((4, 2, 768, "float32", False, True)): ((1, 1), (2, 1), (768, 1)),

  124.     # lenet5

  125.     str((32, 10, 84, 'float16', False, True)): ((1, 1), (16, 1), (84, 1)),

  126.     # alexnet

  127.     str((32, 4096, 9216, 'float32', False, False)): ((1, 1), (16, 1), (1024, 1)),

  128.     str((32, 10, 4096, 'float32', False, False)): ((1, 1), (16, 1), (1024, 1)),

  129. 

  130.     # 3D

  131.     str((128, 128, 64, 1536, "float32", True, False)): ((4, 4), (8, 8), (16, 16), (32, 32)),

  132.     str((128, 768, 128, 64, "float32", False, True)): ((4, 4), (8, 8), (16, 16), (64, 64)),

  133.     str((128, 128, 64, 6144, "float32", True, False)): ((4, 4), (8, 8), (16, 16), (32, 32)),

  134.     str((128, 128, 64, 16384, "float32", True, False)): ((1, 1), (128, 128), (64, 64), (4, 4)),

  135.     str((128, 128, 64, 2048, "float32", True, False)): ((1, 1), (128, 128), (64, 64), (4, 4)),

  136.     str((128, 128, 64, 4096, "float32", True, False)): ((1, 1), (128, 128), (64, 64), (4, 4)),

  137.     str((128, 128, 64, 8192, "float32", True, False)): ((1, 1), (128, 128), (64, 64), (4, 4)),

  138.     str((128, 128, 64, 4096, "float32", True, False)): ((1, 1), (8, 1), (64, 1), (64, 1)),

  139.     str((128, 128, 64, 12288, "float32", True, False)): ((1, 1), (8, 1), (64, 1), (32, 1)),

  140. 

  141.     # 4D

  142.     str((64, 12, 128, 64, 128, "float32", False, False)): ((1, 1), (12, 12), (8, 8), (8, 8), (32, 32)),

  143. 

  144.     str((1, 768, 2, "float32", False, False)): ((768, 1), (1, 1)),

  145.     str((20, 768, 21128, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  146.     str((128, 12, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (64, 1), (128, 1)),

  147.     str((128, 3072, 768, "float32", False, False)): ((1, 1), (768, 1), (32, 1)),

  148.     str((1, 12, 128, 128, 64, "float32", False, True)): ((1, 1), (2, 1), (128, 1), (64, 1)),

  149.     str((128, 128, 64, 12, "float32", True, False)): ((1, 1), (4, 1), (64, 1), (12, 1)),

  150.     str((1, 12, 128, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (64, 1), (32, 1)),

  151.     str((20, 21128, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  152.     str((1, 12, 128, 64, 128, "float32", True, False)): ((1, 1), (8, 1), (64, 1), (32, 1)),

  153.     str((20, 768, 768, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  154.     str((128, 768, 3072, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  155.     str((1, 768, 768, "float32", False, False)): ((768, 1), (16, 1)),

  156.     str((21128, 768, 20, "float32", True, False)): ((8, 1), (768, 1), (8, 1)),

  157.     str((128, 12, 128, 64, "float32", False, True)): ((1, 1), (4, 1), (128, 1), (64, 1)),

  158.     str((128, 768, 768, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  159.     str((40, 768, 768, "float32", False, False)): ((2, 1), (768, 1), (16, 1)),

  160.     str((2, 12, 128, 128, 64, "float32", False, True)): ((1, 1), (1, 1), (2, 1), (128, 1), (64, 1)),

  161.     str((2, 768, 2, "float32", False, False)): ((2, 1), (768, 1), (2, 1)),

  162.     str((40, 21128, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  163.     str((256, 768, 768, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  164.     str((2, 12, 128, 64, 128, "float32", True, False)): ((1, 1), (1, 1), (8, 1), (64, 1), (32, 1)),

  165.     str((2, 12, 128, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (1, 1), (64, 1), (32, 1)),

  166.     str((128, 128, 64, 24, "float32", True, False)): ((1, 1), (2, 1), (64, 1), (12, 1)),

  167.     str((256, 3072, 768, "float32", False, False)): ((1, 1), (768, 1), (32, 1)),

  168.     str((128, 24, 128, 64, "float32", False, True)): ((1, 1), (4, 1), (64, 1), (64, 1)),

  169.     str((2, 768, 768, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  170.     str((21128, 768, 40, "float32", True, False)): ((8, 1), (768, 1), (8, 1)),

  171.     str((40, 768, 21128, "float32", False, False)): ((2, 1), (768, 1), (16, 1)),

  172.     str((256, 768, 3072, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  173.     str((128, 24, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (64, 1), (128, 1)),

  174.     str((128, 48, 128, 64, "float32", False, True)): ((1, 1), (4, 1), (64, 1), (64, 1)),

  175.     str((512, 3072, 768, "float32", False, False)): ((1, 1), (768, 1), (32, 1)),

  176.     str((128, 48, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (64, 1), (128, 1)),

  177.     str((512, 768, 768, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  178.     str((4, 12, 128, 64, 128, "float32", True, False)): ((1, 1), (1, 1), (8, 1), (64, 1), (32, 1)),

  179.     str((512, 768, 3072, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  180.     str((4, 768, 2, "float32", False, False)): ((1, 1), (768, 1), (2, 1)),

  181.     str((80, 768, 768, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  182.     str((4, 12, 128, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (1, 1), (64, 1), (32, 1)),

  183.     str((21128, 768, 80, "float32", True, False)): ((8, 1), (768, 1), (8, 1)),

  184.     str((80, 768, 21128, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  185.     str((128, 128, 64, 48, "float32", True, False)): ((1, 1), (4, 1), (64, 1), (12, 1)),

  186.     str((4, 768, 768, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  187.     str((4, 12, 128, 128, 64, "float32", False, True)): ((1, 1), (1, 1), (2, 1), (128, 1), (64, 1)),

  188.     str((80, 21128, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  189.     str((128, 96, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (64, 1), (128, 1)),

  190.     str((1024, 3072, 768, "float32", False, False)): ((1, 1), (768, 1), (32, 1)),

  191.     str((160, 21128, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  192.     str((8, 12, 128, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (1, 1), (64, 1), (32, 1)),

  193.     str((8, 768, 2, "float32", False, False)): ((8, 1), (768, 1), (2, 1)),

  194.     str((8, 768, 768, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  195.     str((1024, 768, 768, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  196.     str((128, 96, 128, 64, "float32", False, True)): ((1, 1), (4, 1), (64, 1), (64, 1)),

  197.     str((128, 128, 64, 96, "float32", True, False)): ((1, 1), (4, 1), (64, 1), (12, 1)),

  198.     str((8, 12, 128, 128, 64, "float32", False, True)): ((1, 1), (1, 1), (2, 1), (128, 1), (64, 1)),

  199.     str((1024, 768, 3072, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  200.     str((21128, 768, 160, "float32", True, False)): ((8, 1), (768, 1), (8, 1)),

  201.     str((160, 768, 768, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  202.     str((8, 12, 128, 64, 128, "float32", True, False)): ((1, 1), (1, 1), (8, 1), (64, 1), (32, 1)),

  203.     str((160, 768, 21128, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  204.     str((16, 768, 2, "float32", False, False)): ((2, 1), (768, 1), (2, 1)),

  205.     str((320, 21128, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  206.     str((16, 12, 128, 64, 128, "float32", True, False)): ((1, 1), (1, 1), (8, 1), (64, 1), (32, 1)),

  207.     str((2048, 3072, 768, "float32", False, False)): ((1, 1), (768, 1), (32, 1)),

  208.     str((16, 12, 128, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (1, 1), (64, 1), (32, 1)),

  209.     str((128, 128, 64, 192, "float32", True, False)): ((1, 1), (4, 1), (64, 1), (12, 1)),

  210.     str((128, 192, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (64, 1), (128, 1)),

  211.     str((21128, 768, 320, "float32", True, False)): ((8, 1), (768, 1), (8, 1)),

  212.     str((128, 192, 128, 64, "float32", False, True)): ((1, 1), (4, 1), (64, 1), (64, 1)),

  213.     str((320, 768, 21128, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  214.     str((16, 12, 128, 128, 64, "float32", False, True)): ((1, 1), (1, 1), (2, 1), (128, 1), (64, 1)),

  215.     str((320, 768, 768, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  216.     str((2048, 768, 3072, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  217.     str((2048, 768, 768, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  218.     str((16, 768, 768, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  219.     str((4096, 768, 768, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  220.     str((128, 128, 64, 384, "float32", True, False)): ((1, 1), (4, 1), (64, 1), (12, 1)),

  221.     str((4096, 3072, 768, "float32", False, False)): ((1, 1), (768, 1), (32, 1)),

  222.     str((640, 768, 768, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  223.     str((128, 384, 128, 64, "float32", False, True)): ((1, 1), (4, 1), (64, 1), (64, 1)),

  224.     str((128, 384, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (64, 1), (128, 1)),

  225.     str((32, 12, 128, 128, 64, "float32", False, True)): ((1, 1), (1, 1), (2, 1), (128, 1), (64, 1)),

  226.     str((32, 12, 128, 64, 128, "float32", True, False)): ((1, 1), (1, 1), (8, 1), (64, 1), (32, 1)),

  227.     str((32, 12, 128, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (1, 1), (64, 1), (32, 1)),

  228.     str((32, 768, 768, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  229.     str((4096, 768, 3072, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  230.     str((640, 768, 21128, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  231.     str((32, 768, 2, "float32", False, False)): ((4, 1), (768, 1), (2, 1)),

  232.     str((21128, 768, 640, "float32", True, False)): ((8, 1), (768, 1), (8, 1)),

  233.     str((640, 21128, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  234.     str((64, 12, 128, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (1, 1), (64, 1), (32, 1)),

  235.     str((8192, 768, 3072, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  236.     str((1280, 768, 768, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  237.     str((8192, 3072, 768, "float32", False, False)): ((1, 1), (768, 1), (32, 1)),

  238.     str((128, 768, 128, 64, "float32", False, True)): ((1, 1), (4, 1), (64, 1), (64, 1)),

  239.     str((128, 128, 64, 768, "float32", True, False)): ((1, 1), (4, 1), (64, 1), (12, 1)),

  240.     str((21128, 768, 1280, "float32", True, False)): ((8, 1), (768, 1), (8, 1)),

  241.     str((1280, 768, 21128, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  242.     str((64, 12, 128, 64, 128, "float32", True, False)): ((1, 1), (1, 1), (8, 1), (64, 1), (32, 1)),

  243.     str((1280, 21128, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  244.     str((64, 768, 768, "float32", False, False)): ((2, 1), (768, 1), (16, 1)),

  245.     str((64, 12, 128, 128, 64, "float32", False, True)): ((1, 1), (1, 1), (2, 1), (128, 1), (64, 1)),

  246.     str((64, 768, 2, "float32", False, False)): ((8, 1), (768, 1), (2, 1)),

  247.     str((8192, 768, 768, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  248.     str((128, 768, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (64, 1), (128, 1)),

  249.     str((128, 12, 128, 128, 64, "float32", False, True)): ((1, 1), (1, 1), (2, 1), (128, 1), (64, 1)),

  250.     str((16384, 768, 3072, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  251.     str((16384, 768, 768, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  252.     str((128, 1536, 128, 64, "float32", False, True)): ((1, 1), (4, 1), (64, 1), (64, 1)),

  253.     str((128, 768, 2, "float32", False, False)): ((16, 1), (768, 1), (2, 1)),

  254.     str((2560, 768, 768, "float32", False, False)): ((1, 1), (768, 1), (16, 1)),

  255.     str((21128, 768, 2560, "float32", True, False)): ((8, 1), (768, 1), (8, 1)),

  256.     str((128, 12, 128, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (1, 1), (64, 1), (32, 1)),

  257.     str((2560, 768, 21128, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  258.     str((128, 128, 64, 1536, "float32", True, False)): ((1, 1), (4, 1), (64, 1), (12, 1)),

  259.     str((128, 1536, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (64, 1), (128, 1)),

  260.     str((2560, 21128, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  261.     str((128, 12, 128, 64, 128, "float32", True, False)): ((1, 1), (1, 1), (8, 1), (64, 1), (32, 1)),

  262.     str((16384, 3072, 768, "float32", False, False)): ((1, 1), (768, 1), (32, 1)),

  263.     str((1, 16, 128, 64, 128, "float32", False, False)): ((1, 1), (4, 1), (64, 1), (8, 1)),

  264.     str((1, 16, 128, 64, 128, "float32", True, False)): ((1, 1), (8, 1), (64, 1), (4, 1)),

  265.     str((128, 16, 128, 64, "float32", False, True)): ((1, 1), (2, 1), (128, 1), (64, 1)),

  266.     str((128, 1024, 1024, "float32", False, False)): ((2, 1), (1024, 1), (8, 1)),

  267.     str((20, 1024, 21128, "float32", False, False)): ((1, 1), (1024, 1), (8, 1)),

  268.     str((128, 4096, 1024, "float32", False, False)): ((2, 1), (1024, 1), (8, 1)),

  269.     str((128, 1024, 4096, "float32", False, False)): ((2, 1), (1024, 1), (8, 1)),

  270.     str((128, 16, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (64, 1), (32, 1)),

  271.     str((20, 21128, 1024, "float32", False, True)): ((1, 1), (16, 1), (1024, 1)),

  272.     str((1, 16, 128, 128, 64, "float32", False, True)): ((1, 1), (4, 1), (64, 1), (64, 1)),

  273.     str((1, 1024, 2, "float32", False, False)): ((1024, 1), (2, 1)),

  274.     str((21128, 1024, 20, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  275.     str((20, 1024, 1024, "float32", False, False)): ((1, 1), (1024, 1), (8, 1)),

  276.     str((128, 128, 64, 16, "float32", True, False)): ((1, 1), (4, 1), (64, 1), (8, 1)),

  277.     str((1, 1024, 1024, "float32", False, False)): ((1024, 1), (8, 1)),

  278.     str((128, 1024, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  279.     str((1, 1024, 1024, "float32", False, True)): ((8, 1), (1024, 1)),

  280.     str((20, 1024, 1024, "float32", False, True)): ((1, 1), (8, 1), (1024, 1)),

  281.     str((128, 1024, 4096, "float32", False, True)): ((1, 1), (8, 1), (2048, 1)),

  282.     str((128, 4096, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  283.     str((1, 2, 1024, "float32", False, True)): ((2, 1), (1024, 1)),

  284.     str((2, 1024, 1, "float32", True, False)): ((2, 1), (1024, 1)),

  285.     str((1024, 1024, 128, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  286.     str((1024, 1024, 1, "float32", True, False)): ((16, 1), (1024, 1)),

  287.     str((1024, 1024, 20, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  288.     str((4096, 1024, 128, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  289.     str((1024, 4096, 128, "float32", True, False)): ((4, 1), (2048, 1), (2, 1)),

  290.     str((128, 128, 64, 32, "float32", True, False)): ((1, 1), (2, 1), (64, 1), (16, 1)),

  291.     str((128, 32, 128, 64, "float32", False, True)): ((1, 1), (2, 1), (128, 1), (64, 1)),

  292.     str((256, 4096, 1024, "float32", False, False)): ((4, 1), (2048, 1), (2, 1)),

  293.     str((256, 1024, 1024, "float32", False, False)): ((4, 1), (1024, 1), (4, 1)),

  294.     str((40, 21128, 1024, "float32", False, True)): ((1, 1), (16, 1), (1024, 1)),

  295.     str((2, 16, 128, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (2, 1), (64, 1), (16, 1)),

  296.     str((2, 16, 128, 64, 128, "float32", True, False)): ((1, 1), (1, 1), (4, 1), (64, 1), (8, 1)),

  297.     str((21128, 1024, 40, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  298.     str((2, 1024, 2, "float32", False, False)): ((2, 1), (1024, 1), (2, 1)),

  299.     str((256, 1024, 4096, "float32", False, False)): ((2, 1), (1024, 1), (8, 1)),

  300.     str((40, 1024, 21128, "float32", False, False)): ((2, 1), (1024, 1), (8, 1)),

  301.     str((2, 1024, 1024, "float32", False, False)): ((1, 1), (1024, 1), (8, 1)),

  302.     str((40, 1024, 1024, "float32", False, False)): ((2, 1), (1024, 1), (8, 1)),

  303.     str((128, 32, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (64, 1), (32, 1)),

  304.     str((2, 16, 128, 128, 64, "float32", False, True)): ((1, 1), (1, 1), (4, 1), (64, 1), (64, 1)),

  305.     str((2, 1024, 1024, "float32", False, True)): ((1, 1), (8, 1), (1024, 1)),

  306.     str((256, 1024, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  307.     str((256, 4096, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  308.     str((2, 2, 1024, "float32", False, True)): ((1, 1), (2, 1), (1024, 1)),

  309.     str((256, 1024, 4096, "float32", False, True)): ((1, 1), (8, 1), (2048, 1)),

  310.     str((40, 1024, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  311.     str((1024, 1024, 256, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  312.     str((1024, 1024, 40, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  313.     str((4096, 1024, 256, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  314.     str((2, 1024, 2, "float32", True, False)): ((2, 1), (1024, 1), (2, 1)),

  315.     str((1024, 1024, 2, "float32", True, False)): ((16, 1), (1024, 1), (1, 1)),

  316.     str((1024, 4096, 256, "float32", True, False)): ((4, 1), (2048, 1), (2, 1)),

  317.     str((4, 16, 128, 128, 64, "float32", False, True)): ((1, 1), (1, 1), (4, 1), (64, 1), (64, 1)),

  318.     str((4, 16, 128, 64, 128, "float32", True, False)): ((1, 1), (1, 1), (4, 1), (64, 1), (8, 1)),

  319.     str((80, 21128, 1024, "float32", False, True)): ((1, 1), (16, 1), (1024, 1)),

  320.     str((4, 1024, 1024, "float32", False, False)): ((1, 1), (1024, 1), (8, 1)),

  321.     str((128, 64, 128, 64, "float32", False, True)): ((1, 1), (2, 1), (128, 1), (64, 1)),

  322.     str((128, 128, 64, 64, "float32", True, False)): ((1, 1), (2, 1), (64, 1), (16, 1)),

  323.     str((21128, 1024, 80, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  324.     str((4, 1024, 2, "float32", False, False)): ((4, 1), (1024, 1), (2, 1)),

  325.     str((128, 64, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (64, 1), (32, 1)),

  326.     str((80, 1024, 21128, "float32", False, False)): ((1, 1), (1024, 1), (8, 1)),

  327.     str((512, 1024, 1024, "float32", False, False)): ((4, 1), (1024, 1), (4, 1)),

  328.     str((512, 1024, 4096, "float32", False, False)): ((2, 1), (1024, 1), (8, 1)),

  329.     str((512, 4096, 1024, "float32", False, False)): ((4, 1), (2048, 1), (2, 1)),

  330.     str((4, 16, 128, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (2, 1), (64, 1), (16, 1)),

  331.     str((80, 1024, 1024, "float32", False, False)): ((1, 1), (1024, 1), (8, 1)),

  332.     str((512, 4096, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  333.     str((80, 1024, 1024, "float32", False, True)): ((1, 1), (8, 1), (1024, 1)),

  334.     str((4, 1024, 1024, "float32", False, True)): ((1, 1), (8, 1), (1024, 1)),

  335.     str((4, 2, 1024, "float32", False, True)): ((1, 1), (2, 1), (1024, 1)),

  336.     str((512, 1024, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  337.     str((512, 1024, 4096, "float32", False, True)): ((1, 1), (8, 1), (2048, 1)),

  338.     str((1024, 1024, 512, "float32", True, False)): ((4, 1), (2048, 1), (2, 1)),

  339.     str((2, 1024, 4, "float32", True, False)): ((2, 1), (1024, 1), (4, 1)),

  340.     str((4096, 1024, 512, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  341.     str((1024, 1024, 4, "float32", True, False)): ((16, 1), (1024, 1), (1, 1)),

  342.     str((1024, 1024, 80, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  343.     str((1024, 4096, 512, "float32", True, False)): ((4, 1), (2048, 1), (2, 1)),

  344.     str((4096, 768, 3072, "float32", False, True)): ((1, 1), (8, 1), (3072, 1)),

  345.     str((4096, 3072, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  346.     str((4096, 768, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  347.     str((8192, 768, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  348.     str((8192, 3072, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  349.     str((8192, 768, 3072, "float32", False, True)): ((1, 1), (8, 1), (3072, 1)),

  350.     str((16384, 3072, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  351.     str((16384, 768, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  352.     str((16384, 768, 3072, "float32", False, True)): ((1, 1), (8, 1), (3072, 1)),

  353.     str((3072, 768, 128, "float32", True, False)): ((4, 1), (768, 1), (4, 1)),

  354.     str((768, 768, 1, "float32", True, False)): ((16, 1), (768, 1)),

  355.     str((768, 768, 20, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  356.     str((768, 3072, 128, "float32", True, False)): ((4, 1), (3072, 1), (1, 1)),

  357.     str((2, 768, 1, "float32", True, False)): ((2, 1), (768, 1)),

  358.     str((768, 768, 128, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  359.     str((768, 768, 2, "float32", True, False)): ((4, 1), (768, 1), (2, 1)),

  360.     str((768, 768, 40, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  361.     str((3072, 768, 256, "float32", True, False)): ((4, 1), (768, 1), (4, 1)),

  362.     str((768, 3072, 256, "float32", True, False)): ((2, 1), (3072, 1), (2, 1)),

  363.     str((768, 768, 256, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  364.     str((2, 768, 2, "float32", True, False)): ((2, 1), (768, 1), (2, 1)),

  365.     str((768, 768, 80, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  366.     str((3072, 768, 512, "float32", True, False)): ((2, 1), (768, 1), (8, 1)),

  367.     str((768, 3072, 512, "float32", True, False)): ((2, 1), (3072, 1), (2, 1)),

  368.     str((768, 768, 4, "float32", True, False)): ((4, 1), (768, 1), (2, 1)),

  369.     str((768, 768, 512, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  370.     str((2, 768, 4, "float32", True, False)): ((2, 1), (768, 1), (4, 1)),

  371.     str((768, 768, 1024, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  372.     str((768, 3072, 1024, "float32", True, False)): ((2, 1), (3072, 1), (2, 1)),

  373.     str((3072, 768, 1024, "float32", True, False)): ((2, 1), (768, 1), (8, 1)),

  374.     str((768, 768, 8, "float32", True, False)): ((4, 1), (768, 1), (2, 1)),

  375.     str((768, 768, 160, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  376.     str((2, 768, 8, "float32", True, False)): ((2, 1), (768, 1), (8, 1)),

  377.     str((3072, 768, 2048, "float32", True, False)): ((2, 1), (768, 1), (8, 1)),

  378.     str((768, 3072, 2048, "float32", True, False)): ((2, 1), (3072, 1), (2, 1)),

  379.     str((768, 768, 2048, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  380.     str((768, 768, 320, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  381.     str((768, 768, 16, "float32", True, False)): ((4, 1), (768, 1), (2, 1)),

  382.     str((2, 768, 16, "float32", True, False)): ((2, 1), (768, 1), (16, 1)),

  383.     str((768, 768, 32, "float32", True, False)): ((4, 1), (768, 1), (2, 1)),

  384.     str((2, 768, 32, "float32", True, False)): ((2, 1), (768, 1), (16, 1)),

  385.     str((3072, 768, 4096, "float32", True, False)): ((2, 1), (768, 1), (8, 1)),

  386.     str((768, 3072, 4096, "float32", True, False)): ((2, 1), (3072, 1), (2, 1)),

  387.     str((768, 768, 640, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  388.     str((768, 768, 4096, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  389.     str((768, 768, 64, "float32", True, False)): ((4, 1), (768, 1), (2, 1)),

  390.     str((2, 768, 64, "float32", True, False)): ((2, 1), (768, 1), (16, 1)),

  391.     str((768, 3072, 8192, "float32", True, False)): ((2, 1), (3072, 1), (2, 1)),

  392.     str((768, 768, 8192, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  393.     str((768, 768, 1280, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  394.     str((3072, 768, 8192, "float32", True, False)): ((2, 1), (768, 1), (8, 1)),

  395.     str((768, 768, 2560, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  396.     str((3072, 768, 16384, "float32", True, False)): ((2, 1), (768, 1), (8, 1)),

  397.     str((2, 768, 128, "float32", True, False)): ((2, 1), (768, 1), (16, 1)),

  398.     str((768, 768, 16384, "float32", True, False)): ((8, 1), (768, 1), (2, 1)),

  399.     str((768, 3072, 16384, "float32", True, False)): ((2, 1), (3072, 1), (2, 1)),

  400.     str((8, 16, 128, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (2, 1), (64, 1), (16, 1)),

  401.     str((21128, 1024, 160, "float32", True, False)): ((4, 1), (1024, 1), (4, 1)),

  402.     str((1024, 4096, 1024, "float32", False, False)): ((4, 1), (2048, 1), (2, 1)),

  403.     str((160, 1024, 21128, "float32", False, False)): ((2, 1), (1024, 1), (4, 1)),

  404.     str((8, 16, 128, 64, 128, "float32", True, False)): ((1, 1), (1, 1), (4, 1), (64, 1), (8, 1)),

  405.     str((160, 1024, 1024, "float32", False, False)): ((1, 1), (1024, 1), (8, 1)),

  406.     str((8, 1024, 1024, "float32", False, False)): ((1, 1), (1024, 1), (8, 1)),

  407.     str((1024, 1024, 1024, "float32", False, False)): ((4, 1), (1024, 1), (4, 1)),

  408.     str((160, 21128, 1024, "float32", False, True)): ((1, 1), (16, 1), (1024, 1)),

  409.     str((1024, 1024, 4096, "float32", False, False)): ((2, 1), (1024, 1), (8, 1)),

  410.     str((8, 16, 128, 128, 64, "float32", False, True)): ((1, 1), (1, 1), (4, 1), (64, 1), (64, 1)),

  411.     str((128, 128, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (64, 1), (32, 1)),

  412.     str((128, 128, 128, 64, "float32", False, True)): ((1, 1), (2, 1), (128, 1), (64, 1)),

  413.     str((128, 128, 64, 128, "float32", True, False)): ((1, 1), (2, 1), (64, 1), (16, 1)),

  414.     str((8, 1024, 2, "float32", False, False)): ((1, 1), (1024, 1), (1, 1)),

  415.     str((21128, 1024, 320, "float32", True, False)): ((4, 1), (1024, 1), (4, 1)),

  416.     str((16, 1024, 1024, "float32", False, False)): ((1, 1), (1024, 1), (8, 1)),

  417.     str((2048, 4096, 1024, "float32", False, False)): ((4, 1), (2048, 1), (2, 1)),

  418.     str((16, 1024, 2, "float32", False, False)): ((1, 1), (1024, 1), (1, 1)),

  419.     str((16, 16, 128, 64, 128, "float32", True, False)): ((1, 1), (1, 1), (4, 1), (64, 1), (8, 1)),

  420.     str((16, 16, 128, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (2, 1), (64, 1), (16, 1)),

  421.     str((2048, 1024, 1024, "float32", False, False)): ((4, 1), (1024, 1), (4, 1)),

  422.     str((320, 1024, 1024, "float32", False, False)): ((1, 1), (1024, 1), (8, 1)),

  423.     str((2048, 1024, 4096, "float32", False, False)): ((2, 1), (1024, 1), (8, 1)),

  424.     str((320, 21128, 1024, "float32", False, True)): ((1, 1), (16, 1), (1024, 1)),

  425.     str((128, 256, 128, 64, "float32", False, True)): ((1, 1), (2, 1), (128, 1), (64, 1)),

  426.     str((320, 1024, 21128, "float32", False, False)): ((2, 1), (1024, 1), (4, 1)),

  427.     str((128, 256, 64, 128, "float32", False, False)): ((1, 1), (1, 1), (64, 1), (32, 1)),

  428.     str((128, 128, 64, 256, "float32", True, False)): ((1, 1), (2, 1), (64, 1), (16, 1)),

  429.     str((16, 16, 128, 128, 64, "float32", False, True)): ((1, 1), (1, 1), (4, 1), (64, 1), (64, 1)),

  430.     str((1024, 4096, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  431.     str((8, 2, 1024, "float32", False, True)): ((1, 1), (2, 1), (1024, 1)),

  432.     str((160, 1024, 1024, "float32", False, True)): ((1, 1), (8, 1), (1024, 1)),

  433.     str((1024, 1024, 4096, "float32", False, True)): ((1, 1), (8, 1), (2048, 1)),

  434.     str((1024, 1024, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  435.     str((8, 1024, 1024, "float32", False, True)): ((1, 1), (8, 1), (1024, 1)),

  436.     str((16, 2, 1024, "float32", False, True)): ((1, 1), (2, 1), (1024, 1)),

  437.     str((320, 1024, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  438.     str((16, 1024, 1024, "float32", False, True)): ((1, 1), (8, 1), (1024, 1)),

  439.     str((2048, 1024, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  440.     str((2048, 4096, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  441.     str((2048, 1024, 4096, "float32", False, True)): ((1, 1), (8, 1), (2048, 1)),

  442.     str((1024, 4096, 1024, "float32", True, False)): ((4, 1), (2048, 1), (2, 1)),

  443.     str((1024, 1024, 8, "float32", True, False)): ((16, 1), (1024, 1), (1, 1)),

  444.     str((2, 1024, 8, "float32", True, False)): ((2, 1), (1024, 1), (8, 1)),

  445.     str((1024, 1024, 1024, "float32", True, False)): ((4, 1), (2048, 1), (2, 1)),

  446.     str((4096, 1024, 1024, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  447.     str((1024, 1024, 160, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  448.     str((4096, 1024, 2048, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  449.     str((1024, 1024, 320, "float32", True, False)): ((8, 1), (1024, 1), (2, 1)),

  450.     str((1024, 1024, 16, "float32", True, False)): ((16, 1), (1024, 1), (1, 1)),

  451.     str((1024, 4096, 2048, "float32", True, False)): ((4, 1), (2048, 1), (2, 1)),

  452.     str((2, 1024, 16, "float32", True, False)): ((2, 1), (1024, 1), (16, 1)),

  453.     str((1024, 1024, 2048, "float32", True, False)): ((4, 1), (2048, 1), (2, 1)),

  454.     str((32, 1001, 2048, "float16", False, True)): ((1, 1), (77, 1), (256, 1)),

  455.     str((1001, 2048, 32, "float16", True, False)): ((1, 1), (2048, 1), (4, 1)),

  456.     str((32, 2048, 1001, "float16", False, False)): ((1, 1), (2048, 1), (4, 1)),

  457.     str((32, 1001, 2048, "float32", False, True)): ((1, 1), (7, 1), (2048, 1)),

  458.     str((1001, 2048, 32, "float32", True, False)): ((1, 1), (2048, 1), (4, 1)),

  459.     str((32, 2048, 1001, "float32", False, False)): ((1, 1), (2048, 1), (4, 1)),

  460.     str((768, 3072, 131072, "float32", True, False)): ((2, 1), (3072, 1), (2, 1)),

  461.     str((3072, 768, 131072, "float32", True, False)): ((2, 1), (768, 1), (8, 1)),

  462.     str((65536, 768, 3072, "float32", False, True)): ((1, 1), (8, 1), (3072, 1)),

  463.     str((131072, 3072, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  464.     str((32768, 1024, 4096, "float32", False, True)): ((1, 1), (8, 1), (2048, 1)),

  465.     str((65536, 1024, 4096, "float32", False, True)): ((1, 1), (8, 1), (2048, 1)),

  466.     str((131072, 1024, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  467.     str((131072, 768, 3072, "float32", False, True)): ((1, 1), (8, 1), (3072, 1)),

  468.     str((65536, 3072, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  469.     str((32768, 4096, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  470.     str((65536, 4096, 1024, "float32", False, True)): ((2, 1), (8, 1), (1024, 1)),

  471.     str((10240, 21128, 1024, "float32", False, True)): ((1, 1), (16, 1), (1024, 1)),

  472.     str((21128, 1024, 20480, "float32", True, False)): ((4, 1), (1024, 1), (4, 1)),

  473.     str((2048, 3072, 768, "float16", False, False)): ((1, 1), (768, 1), (16, 1)),

  474.     str((2048, 768, 3072, "float16", False, False)): ((2, 1), (768, 1), (8, 1)),

  475.     str((10240, 21128, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  476.     str((20480, 768, 21128, "float32", False, False)): ((2, 1), (768, 1), (8, 1)),

  477.     str((21128, 768, 20480, "float32", True, False)): ((8, 1), (768, 1), (8, 1)),

  478.     str((20480, 21128, 768, "float32", False, True)): ((1, 1), (16, 1), (768, 1)),

  479.     str((32, 10, 2048, "float32", False, True)): ((1, 1), (8, 1), (2048, 1)),

  480.     str((32, 10, 2048, "float16", False, True)): ((1, 1), (8, 1), (2048, 1)),

  481.     str((32, 10, 4096, "float16", False, True)): ((1, 1), (2, 1), (4096, 1)),

  482.     str((768, 768, 16384, "float16", True, False)): ((64, 1), (32, 1), (8, 1)),

  483.     str((3072, 768, 16384, "float16", True, False)): ((64, 1), (32, 1), (8, 1)),

  484.     str((768, 3072, 16384, "float16", True, False)): ((64, 1), (32, 1), (8, 1)),

  485.     str((32, 9216, 4096, "float32", False, False)): ((1, 1), (9216, 1), (1, 1)),

  486.     str((128, 128, 64, 3072, "float32", True, False)): ((1, 1), (64, 1), (64, 1), (1, 1)),

  487.     str((21128, 1024, 10240, "float32", True, False)): ((8, 1), (1024, 1), (1, 1)),

  488.     str((128, 128, 64, 1536, "float16", True, False)): ((1, 1), (1, 1), (16, 1), (1536, 1)),

  489.     str((3072, 768, 2048, "float16", True, False)): ((1, 1), (16, 1), (1024, 1)),

  490.     str((768, 3072, 2048, "float16", True, False)): ((1, 1), (16, 1), (1024, 1)),

  491.     str((4096, 1024, 65536, "float32", True, False)): ((4, 1), (1024, 1), (4, 1)),  # auto tiling crash

  492.     str((1024, 4096, 65536, "float32", True, False)): ((1, 1), (4096, 1), (4, 1)),  # auto tiling crash

  493.     str((16384, 3072, 768, "float16", False, True)): ((1, 1), (16, 1), (768, 1)),  # auto tiling crash

  494.     str((16384, 768, 3072, "float16", False, True)): ((1, 1), (16, 1), (1024, 1)),  # auto tiling crash

  495.     str((1024, 4096, 1024, "float16", False, True)): ((1, 1), (16, 1), (1024, 1)),  # auto tiling crash

  496.     str((131072, 1024, 1024, "float32", False, False)): ((2, 1), (1024, 1), (8, 1)),

  497. 

  498.     # Alexnet shape

  499.     str((32, 10, 4096, "float32", False, True)): ((32, 1), (10, 1), (1, 1)),  # auto tiling crash

  500.     str((32, 4096, 4096, "float16", False, False)): ((1, 1), (16, 1), (512, 1)),  # auto tiling crash

  501.     str((32, 9216, 4096, "float16", False, False)): ((1, 1), (16, 1), (512, 1)),  # auto tiling crash

  502.     str((32, 4096, 9216, "float16", False, True)): ((1, 1), (16, 1), (1024, 1)),  # auto tiling crash

  503.     str((128, 128, 64, 1536, "float16", True, False)): ((1, 1), (1, 1), (16, 1), (96, 1)),  # auto tiling crash

  504.     str((32, 4096, 4096, "float32", False, False)): ((1, 1), (256, 1), (64, 1)),  # performance optimization

  505.     # Alexnet shape

  506. 

  507.     str((768, 3072, 16384, "float16", True, False)): ((64, 1), (32, 1), (8, 1)),

  508. 

  509.     str((3072, 768, 4096, "float16", True, False)): ((2, 1), (768, 1), (8, 1)),

  510.     str((768, 3072, 4096, "float16", True, False)): ((2, 1), (3072, 1), (2, 1)),

  511. 

  512.     str((400, 120, 32, "float16", True, False)): ((2, 1), (32, 1), (2, 1)),

  513. }

  514. 

  515. CORE_NUM = 2 if product_is_mini() else 32

  516. MINIMAL_FOR_MULTICORE = CORE_NUM * 512

  517. 

  518. 

  519. def get_best_align_elem(tensor_size, tensor_dtype):

  520.     """Get the best tiling factor for alignment axis."""

  521.     basic_align_elem = int(ct_util.BLOCK_SIZE / get_bytes(tensor_dtype))

  522.     lcm = least_common_multiple(tensor_size, basic_align_elem)

  523.     gcd = greatest_common_divisor(tensor_size, basic_align_elem)

  524.     if gcd != 1:

  525.         return gcd

  526.     if lcm < tensor_size:

  527.         return min(tensor_size, lcm)

  528.     return -1

  529. 

  530. 

  531. def get_shape_pos_map(tensor_shape):

  532.     """Mapping tensor shape to corresponding axis position."""

  533.     batch_pos = [i for i in range(len(tensor_shape) - 3) if tensor_shape[i] != 1]

  534.     mnk = dict()

  535.     pos_map = {0: "m", 1: "n", 2: "k"}

  536.     count = -1

  537.     for i, _ in enumerate(batch_pos):

  538.         count += 1

  539.         mnk["b%s" % str(i)] = count

  540.     for i, shp in enumerate(tensor_shape[-3:]):

  541.         if shp != 1:

  542.             count += 1

  543.         mnk[pos_map[i]] = count

  544.     return batch_pos, mnk

  545. 

  546. 

  547. def batchmatmul_tiling_strategy(shape, align_dtype, attrs):

  548.     """This is an efficient version of tiling strategy for batchmatmul."""

  549.     if len(shape) < 3:

  550.         raise RuntimeError("Shape must be in the form of [(Batch_out, Batch_in,) M, N, K]. "

  551.                            "Current length of shape is {}".format(len(shape)))

  552.     strategy = list()

  553.     m, n, k = shape[-3:]

  554.     batch_pos, mnk = get_shape_pos_map(shape)

  555.     total_size = reduce(lambda x, y: int(x) * int(y), shape)

  556. 

  557.     # set minimal tile for n as the basic block size for alignment

  558.     align_elem = int(ct_util.BLOCK_SIZE / get_bytes(align_dtype))

  559.     tile_n = get_best_align_elem(n, align_dtype)

  560. 

  561.     # if n is smaller than block size, it is not safe to open multi-core for now

  562.     if n < align_elem:

  563.         n_constraint = [ct_util.TileConstraint.FACTOR]

  564.         used_core = CORE_NUM

  565.         attrs["enable_multicore"] = 0

  566.     else:

  567.         n_constraint = [ct_util.TileConstraint.MOD, ct_util.TileConstraint.MIN]

  568.         used_core = 1

  569. 

  570.     if total_size >= MINIMAL_FOR_MULTICORE and used_core < CORE_NUM:

  571.         # set maximal tile for batch according to multi-core usage

  572.         for i, p in enumerate(batch_pos):

  573.             use = min(shape[p], int((CORE_NUM - 1 + used_core) / used_core))

  574.             used_core *= use

  575.             max_b = int(shape[p] / use)

  576.             strategy.append(ct_util.create_constraint_on_axis(values=max_b,

  577.                                                               constraints=ct_util.TileConstraint.MAX,

  578.                                                               band=0,

  579.                                                               axis=i)[0])

  580. 

  581.     tile_k = get_best_align_elem(k, align_dtype)

  582.     k_constraint = ct_util.TileConstraint.MIN

  583.     total_size /= max(1, int(k / tile_k))

  584. 

  585.     # set minimal tile for m according to multi-core usage when there is no expansion

  586.     tile_m = -1

  587.     m_constraint = ct_util.TileConstraint.MAX

  588.     m_per_block = m

  589.     max_core = min(CORE_NUM, int(total_size / MINIMAL_FOR_MULTICORE))

  590.     if greatest_common_divisor(n, align_elem) != 1 and used_core < max_core:

  591.         left_core = int((max_core - 1 + used_core) / used_core)

  592.         core_limit = max(1, int(m / greatest_common_divisor(left_core, m)))

  593.         nk_in_mem = int(n / max(1, tile_n)) * int(k / tile_k)

  594.         balance_limit = max(1, int(m / greatest_common_divisor(nk_in_mem, m)))

  595.         tile_m = min(core_limit, balance_limit)

  596.         m_per_block = int(m / tile_m)

  597. 

  598.     # for large m case, it is more efficient to balance memory bound and calculation bound

  599.     if m_per_block > int(n / max(1, tile_n)) * int(k / tile_k):

  600.         tile_m = max(min(m, align_elem), tile_m)

  601.         k_constraint = ct_util.TileConstraint.FACTOR

  602. 

  603.     # create constraints based on previous analysis

  604.     if m != 1:

  605.         strategy.append(ct_util.create_constraint_on_axis(values=tile_m,

  606.                                                           constraints=m_constraint,

  607.                                                           band=0,

  608.                                                           axis=mnk["m"])[0])

  609.     if n != 1:

  610.         for constraint in n_constraint:

  611.             strategy.append(ct_util.create_constraint_on_axis(values=tile_n,

  612.                                                               constraints=constraint,

  613.                                                               band=0,

  614.                                                               axis=mnk["n"])[0])

  615. 

  616.     if k != 1:

  617.         strategy.append(ct_util.create_constraint_on_axis(values=tile_k,

  618.                                                           constraints=k_constraint,

  619.                                                           band=0,

  620.                                                           axis=mnk["k"])[0])

  621.     higher_priority_pos = mnk["k"] if k >= n else mnk["n"]

  622.     strategy.append(ct_util.create_constraint_on_axis(values=0,

  623.                                                       constraints=ct_util.TileConstraint.SET_PRIORITY,

  624.                                                       band=0,

  625.                                                       axis=higher_priority_pos)[0])

  626.     strategy.append(ct_util.modify_common_constraints(0.7, ct_util.TileConstraint.SET_MEM_RATIO))

  627.     attrs["custom_tiling"] = strategy

  628.     return attrs

  629. 

  630. def batchmatmul_tiling_strategy_dynamic(shape, output, attrs):

  631.     """This is an efficient version of tiling strategy for batchmatmul."""

  632.     if len(shape) < 3:

  633.         raise RuntimeError("Shape must be in the form of [(Batch_out, Batch_in,) M, N, K]. "

  634.                            "Current length of shape is {}".format(len(shape)))

  635.     strategy = list()

  636.     _, mnk = get_shape_pos_map(shape)

  637.     # create constraints based on previous analysis

  638.     strategy.append(ct_util.create_constraint_on_axis(values=1,

  639.                                                       constraints=ct_util.TileConstraint.FACTOR,

  640.                                                       band=0,

  641.                                                       axis=mnk["m"])[0])

  642.     strategy.append(ct_util.create_constraint_on_axis(values="FULL",

  643.                                                       constraints=ct_util.TileConstraint.MAX,

  644.                                                       band=0,

  645.                                                       axis=mnk["n"])[0])

  646.     strategy.append(ct_util.create_constraint_on_axis(values=8,

  647.                                                       constraints=ct_util.TileConstraint.FACTOR,

  648.                                                       band=0,

  649.                                                       axis=mnk["k"])[0])

  650.     strategy.append(ct_util.modify_common_constraints(0.7, ct_util.TileConstraint.SET_MEM_RATIO))

  651. 

  652.     attrs["custom_tiling"] = strategy

  653.     attrs["dynamic_shape"] = ds.set_dynamic_shape_limit_for_tensor(output, 2048, [1,])

  654.     return attrs

  655. 

  656. def get_mnk_from_matrix(shape_a_list, shape_b_list, trans_a, trans_b):

  657.     """Get m, n and k value from input tensor shapes."""

  658.     m, k = shape_a_list[-2], shape_a_list[-1]

  659.     if trans_a:

  660.         m, k = k, m

  661. 

  662.     n = shape_b_list[-2] if trans_b else shape_b_list[-1]

  663.     return [m, n, k]

  664. 

  665. 

  666. def batchmatmul_set_dim(a_value, b_value, trans_a, trans_b):

  667.     """This function is used to set dim info in attrs by set_dim_map."""

  668.     shape_a_list = get_shape(a_value)

  669.     shape_b_list = get_shape(b_value)

  670.     m, n, k = get_mnk_from_matrix(shape_a_list, shape_b_list, trans_a, trans_b)

  671. 

  672.     key = ()

  673.     if len(shape_a_list) > 2:

  674.         key += tuple(shape_a_list[:-2])

  675. 

  676.     key += (m, n, k, a_value.dtype, trans_a, trans_b)

  677.     set_dims = ct_util.set_dims_by_key(str(key), batchmatmul_set_dim_map)

  678. 

  679.     return set_dims, str(key)

  680. 

  681. 

  682. def batchmatmul_bias_set_dim(a_value, b_value, bias_value, trans_a, trans_b):

  683.     """This function is used to set dim info in attrs by set_dim_map of batchmatmul with bias."""

  684.     return batchmatmul_set_dim(a_value, b_value, trans_a, trans_b)

  685. 

  686. 

  687. def batchmatmul_no_bias_set_dim(a_value, b_value, trans_a, trans_b):

  688.     """This function is used to set dim info in attrs by set_dim_map of batchmatmul without bias."""

  689.     return batchmatmul_set_dim(a_value, b_value, trans_a, trans_b)

  690. 

  691. 

  692. @ct_util.reg_set_dim_func(batchmatmul_bias_set_dim)

  693. @vc_util.check_input_type(akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor, bool, bool)

  694. def batchmatmul_bias(a_value, b_value, bias_value, trans_a, trans_b):

  695.     """

  696.     Multiplies two tensors in batches and adds bias to the output.

  697. 

  698.     Args:

  699.         a_value (tvm.tensor.Tensor): 2D or 3D or 4D tensor of type float16 or float32 with shape(..., r_A, c_A).

  700.         b_value (tvm.tensor.Tensor): 2D or 3D or 4D tensor of same type as a_value with shape(..., r_B, c_B).

  701.         bias_value (tvm.tensor.Tensor): The bias tensor added to the result of a_value * b_value.

  702.                                         Should be of same type as a_value, broadcast is allowed.

  703.         trans_a (bool): Specifies whether a_value is transposed or not, default value is False.

  704.         trans_b (bool): Specifies whether b_value is transposed or not, default value is False.

  705. 

  706.     Returns:

  707.         tvm.tensor.Tensor of same type as a_value with shape(..., r_C, c_C).

  708.             r_C = c_A if trans_a else r_A

  709.             c_C = r_B if trans_b else c_B

  710.     """

  711. 

  712.     if not isinstance(trans_a, bool):

  713.         raise TypeError("trans_a should be of type Boolean.")

  714.     if not isinstance(trans_b, bool):

  715.         raise TypeError("trans_b should be of type Boolean.")

  716.     vc_util.ops_dtype_check([a_value.dtype, b_value.dtype, bias_value.dtype], vc_util.DtypeForDavinci.ALL_FLOAT)

  717.     vc_util.elemwise_dtype_check(a_value.dtype, b_value.dtype)

  718.     vc_util.elemwise_dtype_check(a_value.dtype, bias_value.dtype)

  719.     vc_util.gemm_format_check(get_shape(a_value), get_shape(b_value), trans_a, trans_b)

  720.     if len(a_value.shape) not in [2, 3, 4]:

  721.         raise ValueError("Batch matmul only support 2D, 3D and 4D now.")

  722. 

  723.     c_value = batchmatmul(a_value, b_value, trans_a, trans_b)

  724.     if isinstance(c_value, (tuple, list)):

  725.         c_value = c_value[0]

  726. 

  727.     vc_util.auto_broadcast_check(get_shape(bias_value), get_shape(c_value))

  728. 

  729.     shape_c_list = get_shape(c_value)

  730.     bias_value = akg.topi.broadcast_to(bias_value, shape_c_list)

  731.     dim_info = batchmatmul_bias_set_dim(a_value, b_value, bias_value, trans_a, trans_b)

  732.     if isinstance(dim_info, (tuple, list)):

  733.         dim_info = dim_info[0]

  734.     attrs = {}

  735.     attrs["enable_compute_in_place"] = True

  736.     if dim_info != "":

  737.         attrs["dim"] = dim_info

  738.     batch = get_shape(a_value)[:-2]

  739.     mnk = get_mnk_from_matrix(get_shape(a_value), get_shape(b_value), trans_a, trans_b)

  740.     attrs = batchmatmul_tiling_strategy(batch + mnk, c_value.dtype, attrs)

  741.     return akg.tvm.compute(bias_value.shape,

  742.                            lambda *indice: c_value(*indice) + bias_value(*indice), name='matmul_bias_output'), attrs

  743. 

  744. 

  745. @ct_util.reg_set_dim_func(batchmatmul_no_bias_set_dim)

  746. @vc_util.check_input_type(akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor, bool, bool)

  747. def batchmatmul(a_value, b_value, trans_a=False, trans_b=False):

  748.     """

  749.     Multiplies two tensors in batches.

  750. 

  751.     Args:

  752.         a_value (tvm.tensor.Tensor): 2D or 3D or 4D tensor of type float16 or float32 with shape(..., r_A, c_A).

  753.         b_value (tvm.tensor.Tensor): 2D or 3D or 4D tensor of same type as a_value with shape(..., r_B, c_B).

  754.         trans_a (bool): Specifies whether a_value is transposed or not, default value is False.

  755.         trans_b (bool): Specifies whether b_value is transposed or not, default value is False.

  756. 

  757.     Returns:

  758.         tvm.tensor.Tensor of same type as a_value with shape(..., r_C, c_C).

  759.             r_C = c_A if trans_a else r_A

  760.             c_C = r_B if trans_b else c_B

  761.     """

  762. 

  763.     if not isinstance(trans_a, bool):

  764.         raise TypeError("trans_a should be of type Boolean.")

  765.     if not isinstance(trans_b, bool):

  766.         raise TypeError("trans_b should be of type Boolean.")

  767.     vc_util.ops_dtype_check([a_value.dtype, b_value.dtype], vc_util.DtypeForDavinci.ALL_FLOAT)

  768.     vc_util.elemwise_dtype_check(a_value.dtype, b_value.dtype)

  769.     vc_util.gemm_format_check(get_shape(a_value), get_shape(b_value), trans_a, trans_b)

  770.     if len(a_value.shape) not in [2, 3, 4]:

  771.         raise ValueError("Batch matmul only support 2D, 3D and 4D now.")

  772. 

  773.     dtype = a_value.dtype

  774.     if dtype == 'float16':

  775.         if len(a_value.shape) == 2:

  776.             c_value = vectormatmul_2d_cast(a_value, b_value, trans_a, trans_b, "float32")

  777.         elif len(a_value.shape) == 3:

  778.             c_value = vectormatmul_3d_cast(a_value, b_value, trans_a, trans_b, "float32")

  779.         else:

  780.             c_value = vectormatmul_4d_cast(a_value, b_value, trans_a, trans_b, "float32")

  781. 

  782.     else:

  783.         if len(a_value.shape) == 2:

  784.             c_value = vectormatmul_2d(a_value, b_value, trans_a, trans_b)

  785.         elif len(a_value.shape) == 3:

  786.             c_value = vectormatmul_3d(a_value, b_value, trans_a, trans_b)

  787.         else:

  788.             c_value = vectormatmul_4d(a_value, b_value, trans_a, trans_b)

  789. 

  790.     dim_info = batchmatmul_no_bias_set_dim(a_value, b_value, trans_a, trans_b)

  791.     if isinstance(dim_info, (tuple, list)):

  792.         dim_info = dim_info[0]

  793.     attrs = {}

  794.     attrs["enable_compute_in_place"] = True

  795.     if dim_info != "":

  796.         attrs["dim"] = dim_info

  797. 

  798.     mnk = get_mnk_from_matrix(get_shape(a_value), get_shape(b_value), trans_a, trans_b)

  799.     batch = get_shape(a_value)[:-2]

  800.     is_dynamic = ds.shape_is_dynamic([a_value, b_value])

  801.     if not is_dynamic:

  802.         attrs = batchmatmul_tiling_strategy(batch + mnk, c_value.dtype, attrs)

  803.     else:

  804.         attrs = batchmatmul_tiling_strategy_dynamic(batch + mnk, c_value, attrs)

  805.         attrs["enable_pre_storage_write_simplify"] = True

  806.         attrs["enable_sink_allocate"] = True

  807.         attrs["enable_double_buffer"] = False

  808. 

  809.     return c_value, attrs

  810. 

  811. 

  812. def vectormatmul_3d(a_value, b_value, trans_a, trans_b):

  813.     """hybrid implementation for 3D batchmatmul."""

  814.     if trans_a:

  815.         bs, k, m = a_value.shape

  816.     else:

  817.         bs, m, k = a_value.shape

  818.     if trans_b:

  819.         n = b_value.shape[-2]

  820.     else:

  821.         n = b_value.shape[-1]

  822. 

  823.     dtype = a_value.dtype

  824.     zero = akg.tvm.const(0.0, dtype=dtype)

  825. 

  826.     @script(capture=locals())

  827.     def matmul_hybrid_f_f(a, b, zero):

  828.         t_1 = allocate((bs, m, k, n), a.dtype, 'local')

  829.         t_2 = allocate((bs, m, n), a.dtype, 'local')

  830.         for i_bs in range(0, bs):

  831.             for i_m in range(0, m):

  832.                 for i_k in range(0, k):

  833.                     for i_n in range(0, n):

  834.                         t_1[i_bs, i_m, i_k, i_n] = a[i_bs, i_m, i_k] * b[i_bs, i_k, i_n]

  835.                 for i1_n in range(0, n):

  836.                     t_2[i_bs, i_m, i1_n] = zero

  837.                 for i1_k in range(0, k):

  838.                     for i1_n in range(0, n):

  839.                         t_2[i_bs, i_m, i1_n] = t_2[i_bs, i_m, i1_n] + t_1[i_bs, i_m, i1_k, i1_n]

  840.         return t_2

  841. 

  842.     @script(capture=locals())

  843.     def matmul_hybrid_f_t(a, b, zero):

  844.         t_1 = allocate((bs, m, n, k), a.dtype, 'local')

  845.         t_2 = allocate((bs, m, n), a.dtype, 'local')

  846.         for i_bs in range(0, bs):

  847.             for i_m in range(0, m):

  848.                 for i_n in range(0, n):

  849.                     t_2[i_bs, i_m, i_n] = zero

  850.                     for i_k in range(0, k):

  851.                         t_1[i_bs, i_m, i_n, i_k] = a[i_bs, i_m, i_k] * b[i_bs, i_n, i_k]

  852.                         t_2[i_bs, i_m, i_n] = t_1[i_bs, i_m, i_n, i_k] + t_2[i_bs, i_m, i_n]

  853.         return t_2

  854. 

  855.     @script(capture=locals())

  856.     def matmul_hybrid_t_f(a, b, zero):

  857.         t_1 = allocate((bs, m, k, n), a.dtype, 'local')

  858.         t_2 = allocate((bs, m, n), a.dtype, 'local')

  859.         for i_bs in range(0, bs):

  860.             for i_m in range(0, m):

  861.                 for i_k in range(0, k):

  862.                     for i_n in range(0, n):

  863.                         t_1[i_bs, i_m, i_k, i_n] = a[i_bs, i_k, i_m] * b[i_bs, i_k, i_n]

  864.                 for i1_n in range(0, n):

  865.                     t_2[i_bs, i_m, i1_n] = zero

  866.                 for i1_k in range(0, k):

  867.                     for i1_n in range(0, n):

  868.                         t_2[i_bs, i_m, i1_n] = t_2[i_bs, i_m, i1_n] + t_1[i_bs, i_m, i1_k, i1_n]

  869.         return t_2

  870. 

  871.     if not trans_a and not trans_b:

  872.         c_value = matmul_hybrid_f_f(a_value, b_value, zero)

  873.     elif not trans_a and trans_b:

  874.         c_value = matmul_hybrid_f_t(a_value, b_value, zero)

  875.     elif trans_a and not trans_b:

  876.         c_value = matmul_hybrid_t_f(a_value, b_value, zero)

  877.     else:

  878.         raise ValueError('Not support both transpose yet')

  879. 

  880.     return c_value

  881. 

  882. 

  883. def vectormatmul_4d(a_value, b_value, trans_a, trans_b):

  884.     """hybrid implementation for 4D batchmatmul."""

  885.     if trans_a:

  886.         bs1, bs2, k, m = a_value.shape

  887.     else:

  888.         bs1, bs2, m, k = a_value.shape

  889.     if trans_b:

  890.         n = b_value.shape[-2]

  891.     else:

  892.         n = b_value.shape[-1]

  893. 

  894.     dtype = a_value.dtype

  895.     zero = akg.tvm.const(0.0, dtype=dtype)

  896. 

  897.     @script(capture=locals())

  898.     def matmul_hybrid_f_f(a, b, zero):

  899.         t_1 = allocate((bs1, bs2, m, k, n), a.dtype, 'local')

  900.         t_2 = allocate((bs1, bs2, m, n), a.dtype, 'local')

  901.         for i_bs1 in range(0, bs1):

  902.             for i_bs2 in range(0, bs2):

  903.                 for i_m in range(0, m):

  904.                     for i_k in range(0, k):

  905.                         for i_n in range(0, n):

  906.                             t_1[i_bs1, i_bs2, i_m, i_k, i_n] = a[i_bs1, i_bs2, i_m, i_k] * b[i_bs1, i_bs2, i_k, i_n]

  907.                     for i1_n in range(0, n):

  908.                         t_2[i_bs1, i_bs2, i_m, i1_n] = zero

  909.                     for i1_k in range(0, k):

  910.                         for i1_n in range(0, n):

  911.                             t_2[i_bs1, i_bs2, i_m, i1_n] = t_2[i_bs1, i_bs2, i_m, i1_n] + \

  912.                                 t_1[i_bs1, i_bs2, i_m, i1_k, i1_n]

  913.         return t_2

  914. 

  915.     @script(capture=locals())

  916.     def matmul_hybrid_f_t(a, b, zero):

  917.         t_1 = allocate((bs1, bs2, m, n, k), a.dtype, 'local')

  918.         t_2 = allocate((bs1, bs2, m, n), a.dtype, 'local')

  919.         for i_bs1 in range(0, bs1):

  920.             for i_bs2 in range(0, bs2):

  921.                 for i_m in range(0, m):

  922.                     for i_n in range(0, n):

  923.                         t_2[i_bs1, i_bs2, i_m, i_n] = zero

  924.                         for i_k in range(0, k):

  925.                             t_1[i_bs1, i_bs2, i_m, i_n, i_k] = a[i_bs1, i_bs2, i_m, i_k] * b[i_bs1, i_bs2, i_n, i_k]

  926.                             t_2[i_bs1, i_bs2, i_m, i_n] = t_1[i_bs1, i_bs2, i_m, i_n, i_k] + t_2[i_bs1, i_bs2, i_m, i_n]

  927.         return t_2

  928. 

  929.     @script(capture=locals())

  930.     def matmul_hybrid_t_f(a, b, zero):

  931.         t_1 = allocate((bs1, bs2, m, k, n), a.dtype, 'local')

  932.         t_2 = allocate((bs1, bs2, m, n), a.dtype, 'local')

  933.         for i_bs1 in range(0, bs1):

  934.             for i_bs2 in range(0, bs2):

  935.                 for i_m in range(0, m):

  936.                     for i_k in range(0, k):

  937.                         for i_n in range(0, n):

  938.                             t_1[i_bs1, i_bs2, i_m, i_k, i_n] = a[i_bs1, i_bs2, i_k, i_m] * b[i_bs1, i_bs2, i_k, i_n]

  939.                     for i1_n in range(0, n):

  940.                         t_2[i_bs1, i_bs2, i_m, i1_n] = zero

  941.                     for i1_k in range(0, k):

  942.                         for i1_n in range(0, n):

  943.                             t_2[i_bs1, i_bs2, i_m, i1_n] = t_2[i_bs1, i_bs2, i_m, i1_n] + \

  944.                                 t_1[i_bs1, i_bs2, i_m, i1_k, i1_n]

  945.         return t_2

  946. 

  947.     if not trans_a and not trans_b:

  948.         c_value = matmul_hybrid_f_f(a_value, b_value, zero)

  949.     elif not trans_a and trans_b:

  950.         c_value = matmul_hybrid_f_t(a_value, b_value, zero)

  951.     elif trans_a and not trans_b:

  952.         c_value = matmul_hybrid_t_f(a_value, b_value, zero)

  953.     else:

  954.         raise ValueError('Not support both transpose yet')

  955. 

  956.     return c_value

  957. 

  958. 

  959. def vectormatmul_4d_cast(a_value, b_value, trans_a, trans_b, cast_dtype):

  960.     """dsl implementation with data type cast for 4D batchmatmul."""

  961.     if trans_a:

  962.         b1, b2, k, m = a_value.shape

  963.     else:

  964.         b1, b2, m, k = a_value.shape

  965.     if trans_b:

  966.         n = b_value.shape[-2]

  967.     else:

  968.         n = b_value.shape[-1]

  969. 

  970.     dtype = a_value.dtype

  971. 

  972.     def matmul_4d_dsl(a_value, b_value, trans_a, trans_b):

  973.         if not trans_a and not trans_b:

  974.             ele_mul = akg.tvm.compute((b1, b2, m, n, k),

  975.                                       lambda i_b1, i_b2, i_m, i_n, i_k:

  976.                                       a_value[i_b1, i_b2, i_m, i_k].astype(cast_dtype) *

  977.                                       b_value[i_b1, i_b2, i_k, i_n].astype(cast_dtype),

  978.                                       name="ele_mul")

  979.         elif not trans_a and trans_b:

  980.             ele_mul = akg.tvm.compute((b1, b2, m, n, k),

  981.                                       lambda i_b1, i_b2, i_m, i_n, i_k:

  982.                                       a_value[i_b1, i_b2, i_m, i_k].astype(cast_dtype) *

  983.                                       b_value[i_b1, i_b2, i_n, i_k].astype(cast_dtype),

  984.                                       name="ele_mul")

  985.         elif trans_a and not trans_b:

  986.             ele_mul = akg.tvm.compute((b1, b2, m, n, k),

  987.                                       lambda i_b1, i_b2, i_m, i_n, i_k:

  988.                                       a_value[i_b1, i_b2, i_k, i_m].astype(cast_dtype) *

  989.                                       b_value[i_b1, i_b2, i_k, i_n].astype(cast_dtype),

  990.                                       name="ele_mul")

  991.         elif trans_a and trans_b:

  992.             ele_mul = akg.tvm.compute((b1, b2, m, n, k),

  993.                                       lambda i_b1, i_b2, i_m, i_n, i_k:

  994.                                       b_value[i_b1, i_b2, i_n, i_k].astype(cast_dtype) *

  995.                                       a_value[i_b1, i_b2, i_k, i_m].astype(cast_dtype),

  996.                                       name="ele_mul")

  997.         reduce_axis = akg.tvm.reduce_axis((0, k), name='reduce_axis')

  998.         output_shape = (b1, b2, m, n)

  999.         m_c = akg.tvm.compute(output_shape, lambda *i: akg.tvm.sum(ele_mul[i + (reduce_axis,)],

  1000.                                                                    axis=reduce_axis), name="matmul_compute")

  1001.         return m_c

  1002. 

  1003.     c_cast = matmul_4d_dsl(a_value, b_value, trans_a, trans_b)

  1004.     c_value = cast.cast(c_cast, dtype)

  1005.     if trans_a and trans_b:

  1006.         c_res = akg.topi.transpose(c_value, (1, 0))

  1007.         return c_res

  1008.     return c_value

  1009. 

  1010. 

  1011. def vectormatmul_3d_cast(a_value, b_value, trans_a, trans_b, cast_dtype):

  1012.     """dsl implementation with data type cast for 3D batchmatmul."""

  1013.     if trans_a:

  1014.         b, k, m = a_value.shape

  1015.     else:

  1016.         b, m, k = a_value.shape

  1017.     if trans_b:

  1018.         n = b_value.shape[-2]

  1019.     else:

  1020.         n = b_value.shape[-1]

  1021. 

  1022.     dtype = a_value.dtype

  1023. 

  1024.     def matmul_3d_dsl(a_value, b_value, trans_a, trans_b):

  1025.         if not trans_a and not trans_b:

  1026.             ele_mul = akg.tvm.compute((b, m, n, k), lambda i_b, i_m, i_n, i_k:

  1027.                                       a_value[i_b, i_m, i_k].astype(cast_dtype) *

  1028.                                       b_value[i_b, i_k, i_n].astype(cast_dtype),

  1029.                                       name="ele_mul")

  1030.         elif not trans_a and trans_b:

  1031.             ele_mul = akg.tvm.compute((b, m, n, k), lambda i_b, i_m, i_n, i_k:

  1032.                                       a_value[i_b, i_m, i_k].astype(cast_dtype) *

  1033.                                       b_value[i_b, i_n, i_k].astype(cast_dtype),

  1034.                                       name="ele_mul")

  1035.         elif trans_a and not trans_b:

  1036.             ele_mul = akg.tvm.compute((b, m, n, k), lambda i_b, i_m, i_n, i_k:

  1037.                                       a_value[i_b, i_k, i_m].astype(cast_dtype) *

  1038.                                       b_value[i_b, i_k, i_n].astype(cast_dtype),

  1039.                                       name="ele_mul")

  1040.         elif trans_a and trans_b:

  1041.             ele_mul = akg.tvm.compute((b, m, n, k), lambda i_b, i_m, i_n, i_k:

  1042.                                       b_value[i_b, i_n, i_k].astype(cast_dtype) *

  1043.                                       a_value[i_b, i_k, i_m].astype(cast_dtype),

  1044.                                       name="ele_mul")

  1045.         reduce_axis = akg.tvm.reduce_axis((0, k), name='reduce_axis')

  1046.         output_shape = (b, m, n)

  1047.         m_c = akg.tvm.compute(output_shape, lambda *i: akg.tvm.sum(ele_mul[i + (reduce_axis,)], axis=reduce_axis),

  1048.                               name="matmul_compute")

  1049.         return m_c

  1050. 

  1051.     c_cast = matmul_3d_dsl(a_value, b_value, trans_a, trans_b)

  1052.     c_value = cast.cast(c_cast, dtype)

  1053.     if trans_a and trans_b:

  1054.         c_res = akg.topi.transpose(c_value, (1, 0))

  1055.         return c_res

  1056.     return c_value

  1057. 

  1058. 

  1059. def vectormatmul_2d_cast(a_value, b_value, trans_a, trans_b, cast_dtype):

  1060.     """hybrid implementation with data type cast for 2D batchmatmul."""

  1061.     if trans_a:

  1062.         k, m = a_value.shape

  1063.     else:

  1064.         m, k = a_value.shape

  1065.     if trans_b:

  1066.         n = b_value.shape[-2]

  1067.     else:

  1068.         n = b_value.shape[-1]

  1069. 

  1070.     dtype = a_value.dtype

  1071. 

  1072.     # When the float16 cast to float32 directly, the AutoPoly pass cost a long time.

  1073.     # Therefore, the cast be done in single element.

  1074.     def matmul_2d(a_value, b_value, trans_a, trans_b):

  1075.         if not trans_a and not trans_b:

  1076.             ele_mul = akg.tvm.compute((m, n, k), lambda i_m, i_n, i_k: a_value[i_m, i_k].astype(cast_dtype)

  1077.                                       * b_value[i_k, i_n].astype(cast_dtype),

  1078.                                       name="ele_mul")

  1079.         elif not trans_a and trans_b:

  1080.             ele_mul = akg.tvm.compute((m, n, k), lambda i_m, i_n, i_k: a_value[i_m, i_k].astype(cast_dtype) *

  1081.                                       b_value[i_n, i_k].astype(cast_dtype),

  1082.                                       name="ele_mul")

  1083.         elif trans_a and not trans_b:

  1084.             ele_mul = akg.tvm.compute((m, n, k), lambda i_m, i_n, i_k: a_value[i_k, i_m].astype(cast_dtype) *

  1085.                                       b_value[i_k, i_n].astype(cast_dtype),

  1086.                                       name="ele_mul")

  1087.         elif trans_a and trans_b:

  1088.             ele_mul = akg.tvm.compute((m, n, k), lambda i_m, i_n, i_k: b_value[i_n, i_k].astype(cast_dtype) *

  1089.                                       a_value[i_k, i_m].astype(cast_dtype),

  1090.                                       name="ele_mul")

  1091.         reduce_axis = akg.tvm.reduce_axis((0, k), name='reduce_axis')

  1092.         output_shape = (m, n)

  1093.         m_c = akg.tvm.compute(output_shape, lambda *i: akg.tvm.sum(ele_mul[i + (reduce_axis,)], axis=reduce_axis),

  1094.                               name="matmul_compute")

  1095.         return m_c

  1096. 

  1097.     c_cast = matmul_2d(a_value, b_value, trans_a, trans_b)

  1098.     c_value = cast.cast(c_cast, dtype)

  1099.     if trans_a and trans_b:

  1100.         c_res = akg.topi.transpose(c_value, (1, 0))

  1101.         return c_res

  1102.     return c_value

  1103. 

  1104. 

  1105. def vectormatmul_2d(a_value, b_value, trans_a, trans_b):

  1106.     """hybrid implementation for 2D batchmatmul."""

  1107.     if trans_a:

  1108.         k, m = a_value.shape

  1109.     else:

  1110.         m, k = a_value.shape

  1111.     if trans_b:

  1112.         n = b_value.shape[-2]

  1113.     else:

  1114.         n = b_value.shape[-1]

  1115. 

  1116.     dtype = a_value.dtype

  1117.     zero = akg.tvm.const(0.0, dtype=dtype)

  1118. 

  1119.     @script(capture=locals())

  1120.     def matmul_hybrid_f_f(a, b, zero, mv, nv, kv):

  1121.         t_1 = allocate((mv, kv, nv), a.dtype, 'local')

  1122.         t_2 = output_tensor((mv, nv), a.dtype)

  1123.         for i_m in range(0, mv):

  1124.             for i_k in range(0, kv):

  1125.                 for i_n in range(0, nv):

  1126.                     t_1[i_m, i_k, i_n] = a[i_m, i_k] * b[i_k, i_n]

  1127.             for i1_n in range(0, nv):

  1128.                 t_2[i_m, i1_n] = zero

  1129.             for i1_k in range(0, kv):

  1130.                 for i1_n in range(0, nv):

  1131.                     t_2[i_m, i1_n] = t_2[i_m, i1_n] + t_1[i_m, i1_k, i1_n]

  1132.         return t_2

  1133. 

  1134.     @script(capture=locals())

  1135.     def matmul_hybrid_f_t(a, b, zero, mv, nv, kv):

  1136.         t_1 = allocate((mv, nv, kv), a.dtype, 'local')

  1137.         t_2 = allocate((mv, nv), a.dtype, 'local')

  1138.         for i_m in range(0, mv):

  1139.             for i_n in range(0, nv):

  1140.                 t_2[i_m, i_n] = zero

  1141.                 for i_k in range(0, kv):

  1142.                     t_1[i_m, i_n, i_k] = a[i_m, i_k] * b[i_n, i_k]

  1143.                     t_2[i_m, i_n] = t_1[i_m, i_n, i_k] + t_2[i_m, i_n]

  1144.         return t_2

  1145. 

  1146.     @script(capture=locals())

  1147.     def matmul_hybrid_t_f(a, b, zero, mv, nv, kv):

  1148.         t_1 = allocate((mv, kv, nv), a.dtype, 'local')

  1149.         t_2 = allocate((mv, nv), a.dtype, 'local')

  1150.         for i_m in range(0, mv):

  1151.             for i_k in range(0, kv):

  1152.                 for i_n in range(0, nv):

  1153.                     t_1[i_m, i_k, i_n] = a[i_k, i_m] * b[i_k, i_n]

  1154.             for i1_n in range(0, nv):

  1155.                 t_2[i_m, i1_n] = zero

  1156.             for i1_k in range(0, kv):

  1157.                 for i1_n in range(0, nv):

  1158.                     t_2[i_m, i1_n] = t_2[i_m, i1_n] + t_1[i_m, i1_k, i1_n]

  1159.         return t_2

  1160. 

  1161.     @script(capture=locals())

  1162.     def matmul_hybrid_t_t(a, b, zero, mv, nv, kv):

  1163.         t_1 = allocate((nv, kv, mv), a.dtype, 'local')

  1164.         t_2 = allocate((nv, mv), a.dtype, 'local')

  1165.         for i_n in range(0, nv):

  1166.             for i_m in range(0, mv):

  1167.                 for i_k in range(0, kv):

  1168.                     t_1[i_n, i_k, i_m] = b[i_n, i_k] * a[i_k, i_m]

  1169.             for i1_m in range(0, mv):

  1170.                 t_2[i_n, i1_m] = zero

  1171.             for i1_k in range(0, kv):

  1172.                 for i2_m in range(0, mv):

  1173.                     t_2[i_n, i2_m] = t_2[i_n, i2_m] + t_1[i_n, i1_k, i2_m]

  1174.         return t_2

  1175. 

  1176.     if not trans_a and not trans_b:

  1177.         c_value = matmul_hybrid_f_f(a_value, b_value, zero, m, n, k)

  1178.     elif not trans_a and trans_b:

  1179.         c_value = matmul_hybrid_f_t(a_value, b_value, zero, m, n, k)

  1180.     elif trans_a and not trans_b:

  1181.         c_value = matmul_hybrid_t_f(a_value, b_value, zero, m, n, k)

  1182.     else:

  1183.         c1 = matmul_hybrid_t_t(a_value, b_value, zero, m, n, k)

  1184.         c_value = akg.topi.transpose(c1, (1, 0))

  1185. 

  1186.     return c_value