|
| 1 | +# Copyright 2026, The FlagOS Contributors. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | +# |
| 15 | +# Generated by KernelGen: https://github.com/flagos-ai/KernelGen |
| 16 | +import logging |
| 17 | + |
| 18 | +import torch |
| 19 | +import triton |
| 20 | +import triton.language as tl |
| 21 | + |
| 22 | +from flag_gems import runtime |
| 23 | +from flag_gems.utils import libentry |
| 24 | + |
| 25 | +logger = logging.getLogger(__name__) |
| 26 | + |
| 27 | + |
| 28 | +@libentry() |
| 29 | +@triton.autotune( |
| 30 | + configs=runtime.get_tuned_config("_adaptive_avg_pool3d_backward"), |
| 31 | + key=["n_elements"], |
| 32 | +) |
| 33 | +@triton.jit |
| 34 | +def _adaptive_avg_pool3d_backward_kernel( |
| 35 | + grad_output_ptr, |
| 36 | + grad_input_ptr, |
| 37 | + in_n, |
| 38 | + in_c, |
| 39 | + in_d, |
| 40 | + in_h, |
| 41 | + in_w, |
| 42 | + out_d, |
| 43 | + out_h, |
| 44 | + out_w, |
| 45 | + # Strides for grad_output |
| 46 | + out_stride_n, |
| 47 | + out_stride_c, |
| 48 | + out_stride_d, |
| 49 | + out_stride_h, |
| 50 | + out_stride_w, |
| 51 | + # Strides for grad_input |
| 52 | + grad_in_stride_n, |
| 53 | + grad_in_stride_c, |
| 54 | + grad_in_stride_d, |
| 55 | + grad_in_stride_h, |
| 56 | + grad_in_stride_w, |
| 57 | + n_elements: tl.constexpr, |
| 58 | + BLOCK_SIZE: tl.constexpr, |
| 59 | + MAX_OUT_D: tl.constexpr, |
| 60 | + MAX_OUT_H: tl.constexpr, |
| 61 | + MAX_OUT_W: tl.constexpr, |
| 62 | +): |
| 63 | + pid = tl.program_id(0) |
| 64 | + block_start = pid * BLOCK_SIZE |
| 65 | + offsets = block_start + tl.arange(0, BLOCK_SIZE) |
| 66 | + mask = offsets < n_elements |
| 67 | + |
| 68 | + # Recover 5D coordinates for INPUT from flat index |
| 69 | + n = offsets // (in_c * in_d * in_h * in_w) |
| 70 | + remaining = offsets % (in_c * in_d * in_h * in_w) |
| 71 | + c = remaining // (in_d * in_h * in_w) |
| 72 | + remaining = remaining % (in_d * in_h * in_w) |
| 73 | + d_in = remaining // (in_h * in_w) |
| 74 | + remaining = remaining % (in_h * in_w) |
| 75 | + h_in = remaining // in_w |
| 76 | + w_in = remaining % in_w |
| 77 | + |
| 78 | + # Initialize accumulator |
| 79 | + grad_acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) |
| 80 | + |
| 81 | + # Compute range of possible output positions for each dimension |
| 82 | + d_out_min = (d_in * out_d) // in_d |
| 83 | + d_out_max = ((d_in + 1) * out_d + in_d - 1) // in_d |
| 84 | + h_out_min = (h_in * out_h) // in_h |
| 85 | + h_out_max = ((h_in + 1) * out_h + in_h - 1) // in_h |
| 86 | + w_out_min = (w_in * out_w) // in_w |
| 87 | + w_out_max = ((w_in + 1) * out_w + in_w - 1) // in_w |
| 88 | + |
| 89 | + # Clip to valid range |
| 90 | + d_out_max = tl.minimum(d_out_max, out_d) |
| 91 | + h_out_max = tl.minimum(h_out_max, out_h) |
| 92 | + w_out_max = tl.minimum(w_out_max, out_w) |
| 93 | + |
| 94 | + # Iterate over possible output positions (use static range) |
| 95 | + # Upper bound: each input position can map to at most ceil(out/in) + 1 output |
| 96 | + # positions per dimension. Use a conservative constexpr computed host-side |
| 97 | + # so upsampling (out > in) is handled correctly. |
| 98 | + for d_out in tl.static_range(0, MAX_OUT_D): |
| 99 | + d_out_idx = d_out_min + d_out |
| 100 | + d_out_valid = (d_out_idx >= 0) & (d_out_idx < d_out_max) |
| 101 | + |
| 102 | + for h_out in tl.static_range(0, MAX_OUT_H): |
| 103 | + h_out_idx = h_out_min + h_out |
| 104 | + h_out_valid = (h_out_idx >= 0) & (h_out_idx < h_out_max) |
| 105 | + |
| 106 | + for w_out in tl.static_range(0, MAX_OUT_W): |
| 107 | + w_out_idx = w_out_min + w_out |
| 108 | + w_out_valid = (w_out_idx >= 0) & (w_out_idx < w_out_max) |
| 109 | + |
| 110 | + out_valid = d_out_valid & h_out_valid & w_out_valid |
| 111 | + |
| 112 | + # Compute kernel region for this output |
| 113 | + d_start = (d_out_idx * in_d) // out_d |
| 114 | + d_end = ((d_out_idx + 1) * in_d + out_d - 1) // out_d |
| 115 | + h_start = (h_out_idx * in_h) // out_h |
| 116 | + h_end = ((h_out_idx + 1) * in_h + out_h - 1) // out_h |
| 117 | + w_start = (w_out_idx * in_w) // out_w |
| 118 | + w_end = ((w_out_idx + 1) * in_w + out_w - 1) // out_w |
| 119 | + |
| 120 | + d_end = tl.minimum(d_end, in_d) |
| 121 | + h_end = tl.minimum(h_end, in_h) |
| 122 | + w_end = tl.minimum(w_end, in_w) |
| 123 | + |
| 124 | + # Check if current input is in this output's region |
| 125 | + in_region = ( |
| 126 | + (d_in >= d_start) |
| 127 | + & (d_in < d_end) |
| 128 | + & (h_in >= h_start) |
| 129 | + & (h_in < h_end) |
| 130 | + & (w_in >= w_start) |
| 131 | + & (w_in < w_end) |
| 132 | + ) |
| 133 | + |
| 134 | + # Compute kernel size |
| 135 | + actual_kernel_d = d_end - d_start |
| 136 | + actual_kernel_h = h_end - h_start |
| 137 | + actual_kernel_w = w_end - w_start |
| 138 | + divisor = actual_kernel_d * actual_kernel_h * actual_kernel_w |
| 139 | + |
| 140 | + # Load grad_output and accumulate |
| 141 | + grad_out_ptr = ( |
| 142 | + grad_output_ptr |
| 143 | + + n * out_stride_n |
| 144 | + + c * out_stride_c |
| 145 | + + d_out_idx * out_stride_d |
| 146 | + + h_out_idx * out_stride_h |
| 147 | + + w_out_idx * out_stride_w |
| 148 | + ) |
| 149 | + grad_out_val = tl.load(grad_out_ptr, mask=mask & out_valid) |
| 150 | + |
| 151 | + # Accumulate only if in_region |
| 152 | + contribution = tl.where( |
| 153 | + in_region, grad_out_val / tl.cast(divisor, tl.float32), 0.0 |
| 154 | + ) |
| 155 | + grad_acc += tl.where(out_valid, contribution, 0.0) |
| 156 | + |
| 157 | + # Store result |
| 158 | + grad_in_ptr = ( |
| 159 | + grad_input_ptr |
| 160 | + + n * grad_in_stride_n |
| 161 | + + c * grad_in_stride_c |
| 162 | + + d_in * grad_in_stride_d |
| 163 | + + h_in * grad_in_stride_h |
| 164 | + + w_in * grad_in_stride_w |
| 165 | + ) |
| 166 | + tl.store(grad_in_ptr, grad_acc, mask=mask) |
| 167 | + |
| 168 | + |
| 169 | +def _adaptive_avg_pool3d_backward( |
| 170 | + grad_output: torch.Tensor, |
| 171 | + input: torch.Tensor, |
| 172 | +): |
| 173 | + """Gradient of adaptive_avg_pool3d backward.""" |
| 174 | + logger.debug("GEMS _ADAPTIVE_AVG_POOL3D_BACKWARD") |
| 175 | + |
| 176 | + # Get shapes |
| 177 | + in_n, in_c, in_d, in_h, in_w = input.shape |
| 178 | + out_n, out_c, out_d, out_h, out_w = grad_output.shape |
| 179 | + |
| 180 | + # Allocate output |
| 181 | + grad_input = torch.zeros( |
| 182 | + (in_n, in_c, in_d, in_h, in_w), |
| 183 | + device=input.device, |
| 184 | + dtype=torch.float32, |
| 185 | + ) |
| 186 | + |
| 187 | + if grad_output.numel() == 0: |
| 188 | + return grad_input.to(grad_output.dtype) |
| 189 | + |
| 190 | + n_elements = in_n * in_c * in_d * in_h * in_w |
| 191 | + |
| 192 | + # Upper bound on the number of output positions a single input can map to |
| 193 | + # per dimension. For adaptive pooling, input i contributes to outputs in |
| 194 | + # [o_min, o_max) where o_max - o_min = ceil(out/in) (+1 at boundaries), so a |
| 195 | + # conservative constexpr is ceil(out / in) + 1. This covers upsampling. |
| 196 | + max_out_d = (out_d + in_d - 1) // in_d + 1 |
| 197 | + max_out_h = (out_h + in_h - 1) // in_h + 1 |
| 198 | + max_out_w = (out_w + in_w - 1) // in_w + 1 |
| 199 | + |
| 200 | + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) |
| 201 | + |
| 202 | + _adaptive_avg_pool3d_backward_kernel[grid]( |
| 203 | + grad_output, |
| 204 | + grad_input, |
| 205 | + in_n, |
| 206 | + in_c, |
| 207 | + in_d, |
| 208 | + in_h, |
| 209 | + in_w, |
| 210 | + out_d, |
| 211 | + out_h, |
| 212 | + out_w, |
| 213 | + grad_output.stride(0), |
| 214 | + grad_output.stride(1), |
| 215 | + grad_output.stride(2), |
| 216 | + grad_output.stride(3), |
| 217 | + grad_output.stride(4), |
| 218 | + grad_input.stride(0), |
| 219 | + grad_input.stride(1), |
| 220 | + grad_input.stride(2), |
| 221 | + grad_input.stride(3), |
| 222 | + grad_input.stride(4), |
| 223 | + n_elements, |
| 224 | + MAX_OUT_D=max_out_d, |
| 225 | + MAX_OUT_H=max_out_h, |
| 226 | + MAX_OUT_W=max_out_w, |
| 227 | + ) |
| 228 | + |
| 229 | + return grad_input.to(grad_output.dtype) |
0 commit comments