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

493 lines
23 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: maxpool_ad"""

  18. import akg.tvm

  19. import akg.topi

  20. import akg

  21. from akg.ops.nn import maxpool

  22. from akg.utils.format_transform import get_shape

  23. from akg.utils.dsl_create import cal_pad_shapes_by_strategy

  24. from akg.utils import kernel_exec as utils

  25. from akg.utils import validation_check as vc_util

  26. 

  27. 

  28. @vc_util.check_input_type(akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor, (list, tuple), (list, tuple),

  29.                           (str, list, tuple))

  30. def maxpool_ad_no_custom_diff_poly_all_max(head, data, kernel, stride, pad):

  31.     """automatic differentiate of maxpool with polyhedral"""

  32.     attrs = {"enable_post_poly_loop_partition": False, "enable_pre_poly_loop_partition": False}

  33.     maxpool_fwd = maxpool.old_maxpool(data, kernel, stride, pad)

  34.     [dl_ddata] = akg.differentiate(maxpool_fwd, [data], head, None, None)

  35.     return dl_ddata, attrs

  36. 

  37. 

  38. @vc_util.check_input_type(akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor,

  39.                           (list, tuple), (list, tuple), (str, list, tuple))

  40. def maxpool_ad_no_custom_diff_manual_schedule_all_max(head, data, kernel, stride, pad):

  41.     """automatic differentiate of maxpool with manual schedule."""

  42.     attrs = {"enable_post_poly_loop_partition": False, "enable_pre_poly_loop_partition": False}

  43.     maxpool_fwd = maxpool.old_maxpool(data, kernel, stride, pad)

  44.     [dl_ddata] = akg.differentiate(maxpool_fwd, [data], head, None, None)

  45.     # schedule for differetiation operation

  46.     s = akg.tvm.create_schedule([dl_ddata.op])

  47. 

  48.     new_tensor_red = dl_ddata

  49.     new_tensor = new_tensor_red.op.input_tensors[0]

  50.     data = new_tensor.op.input_tensors[0]

  51.     broadcast = new_tensor.op.input_tensors[1]

  52.     head = new_tensor.op.input_tensors[2]

  53.     forward = broadcast.op.input_tensors[0]

  54. 

  55.     def comp_func(s):

  56.         data_ub = s.cache_read(data, "local.UB", [forward, new_tensor])

  57.         head_ub = s.cache_read(head, "local.UB", [new_tensor])

  58.         result_ub = s.cache_write(new_tensor_red, "local.UB")

  59. 

  60.         s[broadcast].set_scope("local.UB")

  61.         s[forward].set_scope("local.UB")

  62. 

  63.         b, c1, h, w, c0 = forward.op.axis

  64.         oh, ow = forward.op.reduce_axis

  65.         s[forward].reorder(oh, ow, b, c1, h, w, c0)

  66.         s[new_tensor].set_scope("local.UB")

  67. 

  68.         b, c1, h, w, c0 = result_ub.op.axis

  69.         s[result_ub].reorder(*result_ub.op.reduce_axis, b, c1, h, w, c0)

  70. 

  71.         s[broadcast].compute_at(s[result_ub], b)

  72.         s[new_tensor].compute_at(s[result_ub], b)

  73. 

  74.     return dl_ddata, comp_func, attrs

  75. 

  76. 

  77. @vc_util.check_input_type(akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor,

  78.                           akg.tvm.tensor.Tensor, (list, tuple), (list, tuple), (str, list, tuple))

  79. def maxpool_ad(head, data, forward, mask, kernel, stride, pad):

  80.     """automatic differentiate of maxpool with manual schedule."""

  81.     shape = get_shape(data)

  82.     dtype = data.dtype

  83. 

  84.     kernel_h, kernel_w = kernel

  85.     stride_h, stride_w = stride

  86.     [ph_h, _, pw_h, _], [out_size_h, out_size_w] = \

  87.         cal_pad_shapes_by_strategy(shape, kernel, stride, pad)

  88.     batch_size, input_c1, input_h, input_w, input_c0 = shape

  89. 

  90.     # tile size one is proved to be the most efficient one

  91.     tile_scale_h = 1

  92.     tile_scale_w = 1

  93. 

  94.     tile_h = stride_h * tile_scale_h

  95. 

  96.     if kernel_h == stride_h:  # non-overlapping case

  97.         tile_h_pad_u = ph_h % stride_h

  98.     elif kernel_h % stride_h == 0:

  99.         tile_h_pad_u = kernel_h - stride_h - ph_h

  100.     else:

  101.         tile_h_pad_u = kernel_h - kernel_h % stride_h - ph_h

  102.     tile_h_pad_l = kernel_h - stride_h + ph_h

  103.     tile_input_h = tile_h + tile_h_pad_u + tile_h_pad_l

  104.     tile_h_out = (input_h - 1) // tile_h + 1

  105. 

  106.     if ph_h % stride_h == 0:

  107.         pad_output_h = ph_h // stride_h

  108.     else:

  109.         pad_output_h = ph_h // stride_h + 1

  110. 

  111.     if tile_h_pad_u % stride_h == 0:

  112.         pad_output_h -= tile_h_pad_u // stride_h

  113.     else:

  114.         pad_output_h -= tile_h_pad_u // stride_h + 1

  115. 

  116.     tile_output_h = (tile_input_h - kernel_h) // stride_h + 1

  117. 

  118.     tile_w = stride_w * tile_scale_w

  119.     if kernel_w == stride_w:  # non-overlapping case

  120.         tile_w_pad_u = pw_h % stride_w

  121.     elif kernel_w % stride_w == 0:

  122.         tile_w_pad_u = kernel_w - stride_w - pw_h

  123.     else:

  124.         tile_w_pad_u = kernel_w - kernel_w % stride_w - pw_h

  125.     tile_w_pad_l = kernel_w - stride_w + pw_h

  126.     tile_input_w = tile_w + tile_w_pad_u + tile_w_pad_l

  127.     tile_w_out = (input_w - 1) // tile_w + 1

  128. 

  129.     if pw_h % stride_w == 0:

  130.         pad_output_w = pw_h // stride_w

  131.     else:

  132.         pad_output_w = pw_h // stride_w + 1

  133. 

  134.     if tile_w_pad_u % stride_w == 0:

  135.         pad_output_w -= tile_w_pad_u // stride_w

  136.     else:

  137.         pad_output_w -= tile_w_pad_u // stride_w + 1

  138. 

  139.     tile_output_w = (tile_input_w - kernel_w) // stride_w + 1

  140. 

  141.     def custom_maxpool_fdiff(out, inputs, head_, ad_attrs, new_pld_array):

  142.         head_reshaped = akg.tvm.compute((batch_size, input_c1, tile_h_out, tile_w_out,

  143.                                          tile_output_h, tile_output_w, input_c0),

  144.                                         lambda b, c1, h_out, w_out, oh, ow, c0:

  145.                                         akg.tvm.expr.Select(

  146.                                             akg.tvm.any(h_out * tile_scale_h + pad_output_h + oh < 0,

  147.                                                         h_out * tile_scale_h + pad_output_h + oh > out_size_h - 1,

  148.                                                         w_out * tile_scale_w + pad_output_w + ow < 0,

  149.                                                         w_out * tile_scale_w + pad_output_w + ow > out_size_w - 1),

  150.                                             akg.tvm.const(0.0, dtype=dtype),

  151.                                             head_(b, c1,

  152.                                                   h_out * tile_scale_h + pad_output_h + oh,

  153.                                                   w_out * tile_scale_w + pad_output_w + ow,

  154.                                                   c0)),

  155.                                         name="head_reshaped")

  156. 

  157.         mask_reshaped = akg.tvm.compute((batch_size, input_c1, tile_h_out, tile_w_out,

  158.                                          tile_output_h, tile_output_w, kernel_h, kernel_w, input_c0),

  159.                                         lambda b, c1, h_out, w_out, oh, ow, kh, kw, c0:

  160.                                         akg.tvm.expr.Select(

  161.                                             akg.tvm.any(h_out * tile_scale_h + pad_output_h + oh < 0,

  162.                                                         h_out * tile_scale_h + pad_output_h + oh > out_size_h - 1,

  163.                                                         w_out * tile_scale_w + pad_output_w + ow < 0,

  164.                                                         w_out * tile_scale_w + pad_output_w + ow > out_size_w - 1),

  165.                                             akg.tvm.const(0.0, dtype=dtype),

  166.                                             mask(b, c1, kh, kw,

  167.                                                  h_out * tile_scale_h + pad_output_h + oh,

  168.                                                  w_out * tile_scale_w + pad_output_w + ow,

  169.                                                  c0)),

  170.                                         name="mask_reshaped")

  171. 

  172.         d_data = akg.tvm.compute((batch_size, input_c1, tile_h_out, tile_w_out,

  173.                                   tile_output_h, tile_output_w, kernel_h, kernel_w, input_c0),

  174.                                  lambda b, c1, h_out, w_out, oh, ow, kh, kw, c0:

  175.                                  mask_reshaped(b, c1, h_out, w_out, oh, ow, kh, kw, c0)

  176.                                  * head_reshaped(b, c1, h_out, w_out, oh, ow, c0),

  177.                                  name="d_data")

  178. 

  179.         data_reorg = akg.tvm.compute((batch_size, input_c1, tile_h_out, tile_w_out,

  180.                                       tile_output_h, tile_output_w, tile_h, tile_w, input_c0),

  181.                                      lambda b, c1, h_out, w_out, oh, ow, h, w, c0:

  182.                                      akg.tvm.expr.Select(

  183.                                          akg.tvm.any(h + tile_h_pad_u < oh * stride_h,

  184.                                                      h + tile_h_pad_u > oh * stride_h + kernel_h - 1,

  185.                                                      w + tile_w_pad_u < ow * stride_w,

  186.                                                      w + tile_w_pad_u > ow * stride_w + kernel_w - 1),

  187.                                          akg.tvm.const(0, dtype=dtype),

  188.                                          d_data(b, c1, h_out, w_out, oh, ow,

  189.                                                 h + tile_h_pad_u - oh * stride_h,

  190.                                                 w + tile_w_pad_u - ow * stride_w,

  191.                                                 c0)),

  192.                                      name="data_reorg")

  193. 

  194.         result_tile = akg.topi.sum(data_reorg, [4, 5])

  195. 

  196.         result = akg.tvm.compute(shape,

  197.                                  lambda b, c1, h, w, c0:

  198.                                  result_tile(b, c1, h // tile_h, w // tile_w, h % tile_h, w % tile_w, c0),

  199.                                  name="result")

  200.         return [result]

  201. 

  202.     # override differentiation computation with custom function

  203.     [dl_ddata] = akg.differentiate(forward, [data], head, None, None,

  204.                                    override={forward: ([data], custom_maxpool_fdiff)})

  205. 

  206.     # schedule for differetiation operation

  207.     s = akg.tvm.create_schedule([dl_ddata.op])

  208. 

  209.     # get computations

  210.     result = dl_ddata

  211.     result_tile = result.op.input_tensors[0]

  212.     data_reorg = result_tile.op.input_tensors[0]

  213.     d_data = data_reorg.op.input_tensors[0]

  214.     mask_reshaped = d_data.op.input_tensors[0]

  215.     head_reshaped = d_data.op.input_tensors[1]

  216. 

  217.     def comp_func(s):

  218. 

  219.         data_ub = s.cache_read(mask, "local.UB", [mask_reshaped])

  220.         head_ub = s.cache_read(head, "local.UB", [head_reshaped])

  221.         result_ub = s.cache_write(result, "local.UB")

  222. 

  223.         s[d_data].set_scope("local.UB")

  224.         s[data_reorg].set_scope("local.UB")

  225.         s[mask_reshaped].set_scope("local.UB")

  226.         s[head_reshaped].set_scope("local.UB")

  227.         s[result_tile].set_scope("local.UB")

  228. 

  229.         s[result_ub].compute_inline()

  230. 

  231.         # inline inputs

  232.         s[head_ub].compute_inline()

  233.         s[data_ub].compute_inline()

  234. 

  235.         # result_tile dependencies

  236.         s[data_reorg].compute_inline()

  237.         b, c1, h_out, w_out, h, w, c0 = result_tile.op.axis

  238.         oh, ow = result_tile.op.reduce_axis

  239.         s[result_tile].reorder(b, c1, h_out, w_out, h, w, oh, ow, c0)

  240. 

  241.         s[d_data].compute_at(s[result_tile], w_out)

  242.         s[mask_reshaped].compute_at(s[result_tile], w_out)

  243.         s[head_reshaped].compute_at(s[result_tile], w_out)

  244. 

  245.         # tile result

  246.         b, c1, h, w, c0 = result.op.axis

  247.         h_out, h_in = s[result].split(h, tile_h)

  248.         w_out, w_in = s[result].split(w, tile_w)

  249.         s[result].reorder(b, c1, h_out, w_out, h_in, w_in, c0)

  250.         s[result_tile].compute_at(s[result], w_out)

  251. 

  252.     return dl_ddata, comp_func

  253. 

  254. 

  255. @vc_util.check_input_type((list, tuple), (list, tuple), (list, tuple), (str, list, tuple),

  256.                           str, (bool, type(None)), (dict, type(None)))

  257. def maxpool_ad_manual_schedule_all_max(shape, kernel, stride, pad, dtype, polyhedral=True, attrs=None):

  258.     """automatic differentiate of maxpool with manual schedule for all maximum."""

  259.     kernel_h, kernel_w = kernel

  260.     stride_h, stride_w = stride

  261.     pad_h, pad_w, _, _ = pad

  262.     batch_size, input_c1, input_h, input_w, input_c0 = shape

  263.     pad_shape = (batch_size, input_c1, input_h + 2 * pad_h, input_w + 2 * pad_w, input_c0)

  264.     out_size_h = (input_h + 2 * pad_h - kernel_h) // stride_h + 1

  265.     out_size_w = (input_w + 2 * pad_w - kernel_w) // stride_w + 1

  266.     out_shape = (batch_size, input_c1, out_size_h, out_size_w, input_c0)

  267. 

  268.     def custom_maxpool_fdiff(out, inputs, head_, ad_attrs, new_pld_array):

  269.         in_data = inputs[0]

  270. 

  271.         data_separated_by_windows = (kernel_h, kernel_w, batch_size, input_c1, out_size_h, out_size_w, input_c0)

  272. 

  273.         pad_data = akg.tvm.compute(pad_shape,

  274.                                    lambda b, c1, h, w, c0:

  275.                                    akg.tvm.expr.Select(

  276.                                        akg.tvm.all(h >= pad_h,

  277.                                                    h < input_h + pad_h,

  278.                                                    w >= pad_w,

  279.                                                    w < input_w + pad_w),

  280.                                        in_data(b, c1, h - pad_h, w - pad_w, c0),

  281.                                        akg.tvm.const(0.0, dtype=dtype)),

  282.                                    name="pad_data")

  283. 

  284.         data_reshaped = akg.tvm.compute(data_separated_by_windows,

  285.                                         lambda wh, ww, b, c1, oh, ow, c0:

  286.                                         pad_data(b, c1, oh * stride_h + wh, ow * stride_w + ww, c0),

  287.                                         name="data_reshaped")

  288. 

  289.         max_broadcast = akg.tvm.compute(data_separated_by_windows,

  290.                                         lambda wh, ww, b, c1, oh, ow, c0:

  291.                                         out(b, c1, oh, ow, c0),

  292.                                         name="max_broadcast")

  293. 

  294.         equal = akg.tvm.compute(data_separated_by_windows,

  295.                                 lambda wh, ww, b, c1, oh, ow, c0:

  296.                                 akg.tvm.expr.Select(

  297.                                     max_broadcast(wh, ww, b, c1, oh, ow, c0) ==

  298.                                     data_reshaped(wh, ww, b, c1, oh, ow, c0),

  299.                                     head_(b, c1, oh, ow, c0),

  300.                                     akg.tvm.const(0.0, dtype=dtype)),

  301.                                 name="equal")

  302. 

  303.         data_reorg = akg.tvm.compute((out_size_h, out_size_w, batch_size, input_c1, input_h + 2 * pad_h,

  304.                                       input_w + 2 * pad_w, input_c0),

  305.                                      lambda oh, ow, b, c1, h, w, c0:

  306.                                      akg.tvm.expr.Select(

  307.                                          akg.tvm.any(h < oh * stride_h,

  308.                                                      h > oh * stride_h + kernel_h - 1,

  309.                                                      w < ow * stride_w,

  310.                                                      w > ow * stride_w + kernel_w - 1),

  311.                                          akg.tvm.const(0, dtype=dtype),

  312.                                          equal(h - oh * stride_h, w - ow * stride_w, b, c1, oh, ow, c0)),

  313.                                      name="data_reorg")

  314. 

  315.         result_pad = akg.topi.sum(data_reorg, [0, 1])

  316. 

  317.         result = akg.tvm.compute(shape,

  318.                                  lambda b, c1, h, w, c0:

  319.                                  result_pad(b, c1, h + pad_h, w + pad_w, c0),

  320.                                  name="result")

  321. 

  322.         return [result]

  323. 

  324.     # tensor for the input data

  325.     data = akg.tvm.placeholder(shape, dtype, name="input_data")

  326. 

  327.     # maxpool output

  328.     forward = akg.tvm.placeholder(out_shape, name="forward", dtype=dtype)

  329. 

  330.     # adjoint tensor for the differentiation

  331.     head = akg.tvm.placeholder(out_shape, name="head", dtype=dtype)

  332. 

  333.     # override differentiation computation with custom function

  334.     [dl_ddata] = akg.differentiate(forward, [data], head, None, None,

  335.                                    override={forward: ([data], custom_maxpool_fdiff)})

  336. 

  337.     # schedule for differetiation operation

  338.     s = akg.tvm.create_schedule([dl_ddata.op])

  339. 

  340.     # get computations

  341.     result = dl_ddata

  342.     result_pad = result.op.input_tensors[0]

  343.     data_reorg = result_pad.op.input_tensors[0]

  344.     equal = data_reorg.op.input_tensors[0]

  345.     max_broadcast = equal.op.input_tensors[0]

  346.     data_reshaped = equal.op.input_tensors[1]

  347.     pad_data = data_reshaped.op.input_tensors[0]

  348. 

  349.     data_ub = s.cache_read(data, "local.UB", [pad_data])

  350.     head_ub = s.cache_read(head, "local.UB", [equal])

  351.     forward_ub = s.cache_read(forward, "local.UB", [max_broadcast])

  352.     result_ub = s.cache_write(result, "local.UB")

  353. 

  354.     s[max_broadcast].set_scope("local.UB")

  355.     s[data_reshaped].set_scope("local.UB")

  356.     s[pad_data].set_scope("local.UB")

  357.     s[equal].set_scope("local.UB")

  358.     s[data_reorg].set_scope("local.UB")

  359.     s[result_pad].set_scope("local.UB")

  360. 

  361.     s[data_ub].compute_inline()

  362.     s[result_ub].compute_inline()

  363.     s[pad_data].compute_inline()

  364. 

  365.     # equal dependencies

  366.     s[forward_ub].compute_at(s[equal], equal.op.axis[0])

  367.     s[max_broadcast].compute_at(s[equal], equal.op.axis[0])

  368.     s[data_reshaped].compute_at(s[equal], equal.op.axis[0])

  369.     s[head_ub].compute_at(s[equal], equal.op.axis[0])

  370. 

  371.     s[equal].compute_at(s[result_pad], result_pad.op.axis[0])

  372. 

  373.     # result dependencies

  374.     s[data_reorg].compute_inline()

  375.     b, c1, h, w, c0 = result_pad.op.axis

  376.     oh, ow = result_pad.op.reduce_axis

  377.     s[result_pad].reorder(oh, ow, b, c1, h, w, c0)

  378.     # s[result_pad].compute_at(s[result], result.op.axis[1])

  379. 

  380.     b, c1, h, w, c0 = result.op.axis

  381.     h_out, _ = s[result].split(h, stride_h)

  382.     s[result_pad].compute_at(s[result], h_out)

  383. 

  384.     with akg.build_config(add_lower_pass=utils.debug_mode(0), dump_pass_ir=True):

  385.         mod = akg.build(s, [head, data, forward, dl_ddata], "cce", name="maxpool_ad_manual_schedule_all_max",

  386.                         attrs=attrs, polyhedral=polyhedral)

  387.         source_code = mod.imported_modules[0].get_source()

  388.         kernel_name = "maxpool_ad_manual_schedule_all_max"

  389.         utils.create_code(kernel_name, './', source_code)

  390.     return mod

  391. 

  392. 

  393. def maxpool_ad_manual_schedule_no_overlap_all_max(shape, kernel, stride, pad, dtype, attrs=None, polyhedral=False):

  394.     """automatic differentiate of maxpool with manual schedule for no overlap case."""

  395.     kernel_h, kernel_w = kernel

  396.     stride_h, stride_w = stride

  397.     pad_h, pad_w, _, _ = pad

  398.     batch_size, input_c1, input_h, input_w, input_c0 = shape

  399.     pad_shape = (batch_size, input_c1, input_h + 2 * pad_h, input_w + 2 * pad_w, input_c0)

  400. 

  401.     def custom_maxpool_fdiff(out, inputs, head_, ad_attrs, new_pld_array):

  402.         in_data = inputs[0]

  403. 

  404.         if stride_w != kernel_w:

  405.             raise RuntimeError("Only supports kernels with same dimensions as stride size!")

  406.         if stride_h != kernel_h:

  407.             raise RuntimeError("Only supports kernels with same dimensions as stride size!")

  408. 

  409.         out_broadcast = akg.tvm.compute(pad_shape,

  410.                                         lambda b, c1, h, w, c0:

  411.                                         out(b, c1, akg.tvm.floordiv(h, stride_h), akg.tvm.floordiv(w, stride_w), c0),

  412.                                         name="out_broadcast")

  413. 

  414.         # copy output to the shape of the padded input, copying the same value for the entire kernel size

  415.         out_broadcast = akg.tvm.compute(pad_shape,

  416.                                         lambda b, c1, h, w, c0:

  417.                                         out(b, c1, akg.tvm.floordiv(h, stride_h), akg.tvm.floordiv(w, stride_w), c0),

  418.                                         name="out_broadcast")

  419. 

  420.         # copy head to the shape of the padded input, copying the same value for the entire kernel size

  421.         head_broadcast = akg.tvm.compute(pad_shape,

  422.                                          lambda b, c1, h, w, c0:

  423.                                          head_(b, c1, akg.tvm.floordiv(h, stride_h), akg.tvm.floordiv(w, stride_w), c0),

  424.                                          name="head_broadcast")

  425. 

  426.         # check if value was a maximum and assign head of that position if it was

  427.         # this is done for all the maximum values within one kernel

  428.         result = akg.tvm.compute(in_data.shape,

  429.                                  lambda b, c1, h, w, c0:

  430.                                  akg.tvm.expr.Select(

  431.                                      in_data(b, c1, h, w, c0) == out_broadcast(b, c1, h + pad_h, w + pad_w, c0),

  432.                                      head_broadcast(b, c1, h + pad_h, w + pad_w, c0),

  433.                                      akg.tvm.const(0, dtype=in_data.dtype)),

  434.                                  name="result")

  435.         return [result]

  436. 

  437.     out_size_h = (input_h + 2 * pad_h - kernel_h) // stride_h + 1

  438.     out_size_w = (input_w + 2 * pad_w - kernel_w) // stride_w + 1

  439. 

  440.     out_shape = (batch_size, input_c1, out_size_h, out_size_w, input_c0)

  441. 

  442.     # tensor for the input data

  443.     data = akg.tvm.placeholder(shape, dtype, name="input_data")

  444. 

  445.     # maxpool output

  446.     forward = akg.tvm.placeholder(out_shape, name="forward", dtype=dtype)

  447. 

  448.     # adjoint tensor for the differentiation

  449.     head = akg.tvm.placeholder(out_shape, name="head", dtype=dtype)

  450. 

  451.     # override differentiation computation with custom function

  452.     [dl_ddata] = akg.differentiate(forward, [data], head, None, None,

  453.                                    override={forward: ([data], custom_maxpool_fdiff)})

  454. 

  455.     # schedule for differetiation operation

  456.     s = akg.tvm.create_schedule([dl_ddata.op])

  457. 

  458.     # get computations

  459.     result = dl_ddata

  460.     forward_broadcast = result.op.input_tensors[1]

  461.     head_broadcast = result.op.input_tensors[2]

  462. 

  463.     # cache reads and writes

  464.     result_ub = s.cache_write(result, "local.UB")

  465.     data_ub = s.cache_read(data, "local.UB", [result_ub])

  466.     head_ub = s.cache_read(head, "local.UB", [head_broadcast])

  467.     forward_ub = s.cache_read(forward, "local.UB", [forward_broadcast])

  468. 

  469.     s[head_broadcast].set_scope("local.UB")

  470.     s[forward_broadcast].set_scope("local.UB")

  471. 

  472.     s[head_ub].compute_at(s[head_broadcast], head_broadcast.op.axis[0])

  473.     s[forward_ub].compute_at(s[forward_broadcast], forward_broadcast.op.axis[0])

  474.     s[data_ub].compute_at(s[result_ub], result_ub.op.axis[0])

  475.     s[forward_broadcast].compute_at(s[result_ub], result_ub.op.axis[0])

  476.     s[head_broadcast].compute_at(s[result_ub], result_ub.op.axis[0])

  477. 

  478.     _, c1, h, _, _ = result.op.axis

  479. 

  480.     if input_h + 2 * pad_h > 32 or input_w + 2 * pad_w > 32:

  481.         h_outer, _ = s[result].split(h, 4)

  482.         s[result_ub].compute_at(s[result], h_outer)

  483.     else:

  484.         s[result_ub].compute_at(s[result], c1)

  485. 

  486.     with akg.build_config(add_lower_pass=utils.debug_mode(0), dump_pass_ir=True):

  487.         mod = akg.build(s, [head, data, forward, dl_ddata], "cce",

  488.                         name="maxpool_ad_manual_schedule_no_overlap_all_max", attrs=attrs, polyhedral=polyhedral)

  489.         source_code = mod.imported_modules[0].get_source()

  490.         kernel_name = "maxpool_ad_manual_schedule_no_overlap_all_max"

  491.         utils.create_code(kernel_name, './', source_code)

  492.     return mod