已合并
[Inductor] bugfix/rms norm simd multi reduction #39499
luqichao创建于 6月29日
[Inductor] bugfix/rms norm simd multi reduction #39499
已合并
共 3 个文件变更+156-27
| @@ -25,6 +25,23 @@ class TestReduction(TestUtils): | |||
| 25 | view: "f32[9600, 2304]" = torch.ops.aten.view.default(add_3, [9600, 2304]) | 25 | view: "f32[9600, 2304]" = torch.ops.aten.view.default(add_3, [9600, 2304]) |
| 26 | return [None, primals_5, getitem_3, rsqrt, add_2, view, primals_2] | 26 | return [None, primals_5, getitem_3, rsqrt, add_2, view, primals_2] |
| 27 | 27 | ||
| 28 | + def rms_norm_weight_grad(self, grad_out_base, q, q_square_sum, permute_order): | ||
| 29 | + grad_out = grad_out_base.permute(*permute_order) | ||
| 30 | + inv_rms = torch.rsqrt(q_square_sum.unsqueeze(-1) / q.shape[-1] + 1e-6) | ||
| 31 | + grad_weight = (grad_out * q.float() * inv_rms).sum(dim=(0, 1, 2)) | ||
| 32 | + return grad_weight.to(torch.bfloat16) | ||
| 33 | + | ||
| 34 | + def check_rms_norm_weight_grad(self, grad_out_base, q, q_square_sum, permute_order): | ||
| 35 | + expected = self.rms_norm_weight_grad( | ||
| 36 | + grad_out_base, q, q_square_sum, permute_order | ||
| 37 | + ) | ||
| 38 | + compiled = torch.compile( | ||
| 39 | + self.rms_norm_weight_grad, backend="inductor", dynamic=False | ||
| 40 | + ) | ||
| 41 | + actual = compiled(grad_out_base, q, q_square_sum, permute_order) | ||
| 42 | + | ||
| 43 | + self.assertEqual(expected, actual, atol=1e-3, rtol=1e-3) | ||
| 44 | + | ||
| 28 | def test_reduction_cases_shapes(self): | 45 | def test_reduction_cases_shapes(self): |
| 29 | device = 'npu' | 46 | device = 'npu' |
| 30 | primals_2: "f32[32, 2304]" = torch.randn((32, 2304), device=device, dtype=torch.float32) | 47 | primals_2: "f32[32, 2304]" = torch.randn((32, 2304), device=device, dtype=torch.float32) |
| @@ -43,6 +60,22 @@ class TestReduction(TestUtils): | |||
| 43 | self.assertEqual(view_ref, view, atol=1e-3, rtol=1e-3, equal_nan=True) | 60 | self.assertEqual(view_ref, view, atol=1e-3, rtol=1e-3, equal_nan=True) |
| 44 | self.assertEqual(primals_2_ref, primals_2, atol=1e-3, rtol=1e-3, equal_nan=True) | 61 | self.assertEqual(primals_2_ref, primals_2, atol=1e-3, rtol=1e-3, equal_nan=True) |
| 45 | 62 | ||
| 63 | + def test_rms_norm_weight_grad_head_seq_permute(self): | ||
| 64 | + device = "npu" | ||
| 65 | + grad_out_base = torch.randn((2, 3, 4, 8), device=device, dtype=torch.bfloat16) | ||
| 66 | + q = torch.randn((2, 4, 3, 8), device=device, dtype=torch.bfloat16) | ||
| 67 | + q_square_sum = torch.rand((2, 4, 3), device=device, dtype=torch.float32) | ||
| 68 | + | ||
| 69 | + self.check_rms_norm_weight_grad(grad_out_base, q, q_square_sum, (0, 2, 1, 3)) | ||
| 70 | + | ||
| 71 | + def test_rms_norm_weight_grad_batch_seq_permute(self): | ||
| 72 | + device = "npu" | ||
| 73 | + grad_out_base = torch.randn((4, 2, 3, 8), device=device, dtype=torch.bfloat16) | ||
| 74 | + q = torch.randn((2, 4, 3, 8), device=device, dtype=torch.bfloat16) | ||
| 75 | + q_square_sum = torch.rand((2, 4, 3), device=device, dtype=torch.float32) | ||
| 76 | + | ||
| 77 | + self.check_rms_norm_weight_grad(grad_out_base, q, q_square_sum, (1, 0, 2, 3)) | ||
| 78 | + | ||
| 46 | 79 | ||
| 47 | if __name__ == "__main__": | 80 | if __name__ == "__main__": |
| 48 | run_tests() | 81 | run_tests() |
| @@ -364,6 +364,11 @@ class ReductionAnalysis: | |||
| 364 | if not reduction_layout_var_list: | 364 | if not reduction_layout_var_list: |
| 365 | raise RuntimeError("assert reduction_layout_var_list is not empty") | 365 | raise RuntimeError("assert reduction_layout_var_list is not empty") |
| 366 | 366 | ||
| 367 | + if self.numof_reduction_axis() > 1 and self.contiguous_reduction: | ||
X | |||
| 368 | + if not self.kernel.golden_var_list: | ||
| 369 | + self.kernel.select_golden_varlist() | ||
| 370 | + return sum(1 for x in self.kernel.golden_var_list if x.name[0] != 'r') | ||
| 371 | + | ||
| 367 | dim = -1 | 372 | dim = -1 |
| 368 | for i, x in enumerate(reversed(reduction_layout_var_list)): | 373 | for i, x in enumerate(reversed(reduction_layout_var_list)): |
| 369 | if x.name[0] == 'r': | 374 | if x.name[0] == 'r': |
| @@ -2124,12 +2124,15 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2124 | if not (self.loads or self.stores or self.compute or self.post_loop_store): | 2124 | if not (self.loads or self.stores or self.compute or self.post_loop_store): |
| 2125 | return | 2125 | return |
| 2126 | 2126 | ||
| 2127 | - def write_pointwise(): | 2127 | + def write_pointwise(allow_stores=None): |
| 2128 | + if allow_stores is None: | ||
| 2129 | + allow_stores = self.numof_reduction_axis() <= 1 | ||
| 2128 | self._emit_coordinate_transforms() | 2130 | self._emit_coordinate_transforms() |
| 2129 | self.body.splice(self.indexing_code) | 2131 | self.body.splice(self.indexing_code) |
| 2130 | self.body.splice(self.loads) | 2132 | self.body.splice(self.loads) |
| 2131 | self.body.splice(self.compute) | 2133 | self.body.splice(self.compute) |
| 2132 | - self.body.splice(self.stores) | 2134 | + if allow_stores: |
| 2135 | + self.body.splice(self.stores) | ||
| 2133 | 2136 | ||
| 2134 | def collect_store_unified_vars(): | 2137 | def collect_store_unified_vars(): |
| 2135 | """ | 2138 | """ |
| @@ -2243,6 +2246,23 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2243 | 2246 | ||
| 2244 | reduction_1d = is_1d_reduction() | 2247 | reduction_1d = is_1d_reduction() |
| 2245 | do_indent = False | 2248 | do_indent = False |
| 2249 | + | ||
| 2250 | + have_load_store = self.find_axis_in_load_store(range_val) | ||
| 2251 | + if not have_load_store: | ||
| 2252 | + indexing_code = None | ||
| 2253 | + | ||
| 2254 | + is_first_reduction_tiling = ( | ||
| 2255 | + self.numof_reduction_axis() > 1 | ||
| 2256 | + and range_val.is_tiling_axis | ||
| 2257 | + and range_val.prefix == "r" | ||
| 2258 | + and not any(ax.prefix == "r" for ax in self.sorted_axis[:index]) | ||
| 2259 | + ) | ||
| 2260 | + use_outer_reduction_post_loop = ( | ||
| 2261 | + self.numof_reduction_axis() > 1 | ||
| 2262 | + and range_val.prefix == "r" | ||
| 2263 | + and bool(self.prefix._lines) | ||
| 2264 | + ) | ||
| 2265 | + | ||
| 2246 | # tiling axis and last tiling | 2266 | # tiling axis and last tiling |
| 2247 | if range_val.is_tiling_axis and last_tiling: | 2267 | if range_val.is_tiling_axis and last_tiling: |
| 2248 | do_indent = False | 2268 | do_indent = False |
| @@ -2253,23 +2273,31 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2253 | if ( | 2273 | if ( |
| 2254 | range_val.prefix != "r" or not self.persistent_reduction | 2274 | range_val.prefix != "r" or not self.persistent_reduction |
| 2255 | ) and need_axis_loop: | 2275 | ) and need_axis_loop: |
| 2256 | - self.body.splice(self.prefix) | 2276 | + if self.numof_reduction_axis() <= 1: |
| 2277 | + self.body.splice(self.prefix) | ||
| 2257 | self.body.writeline( | 2278 | self.body.writeline( |
| 2258 | f"for loop_{range_val.name} in range(loops_{range_val.name}):" | 2279 | f"for loop_{range_val.name} in range(loops_{range_val.name}):" |
| 2259 | ) | 2280 | ) |
| 2260 | do_indent = True | 2281 | do_indent = True |
| 2261 | loop_body(index, indexing_code, is_last_axis, do_indent) | 2282 | loop_body(index, indexing_code, is_last_axis, do_indent) |
| 2262 | - self.body.splice(self.post_loop_combine) | 2283 | + if use_outer_reduction_post_loop: |
| 2263 | - self.body.splice(self.post_loop_store) | 2284 | + pass |
| 2264 | - # Output deferred reduction stores here (outside the loop). | 2285 | + else: |
| 2265 | - # body.writeline() controls indentation uniformly, keeping | 2286 | + if self.numof_reduction_axis() <= 1 or range_val.prefix != "r": |
| 2266 | - # these stores consistent with post_loop_store in both | 2287 | + self.body.splice(self.post_loop_combine) |
| 2267 | - # static and dynamic modes. | 2288 | + self.body.splice(self.post_loop_store) |
| 2268 | - for store_line in self._deferred_reduction_stores: | 2289 | + # Output deferred reduction stores here (outside the loop). |
| 2269 | - self.body.writeline(store_line) | 2290 | + # body.writeline() controls indentation uniformly, keeping |
| 2270 | - self._deferred_reduction_stores.clear() | 2291 | + # these stores consistent with post_loop_store in both |
| 2271 | - self.post_loop_combine.clear() | 2292 | + # static and dynamic modes. |
| 2272 | - self.post_loop_store.clear() | 2293 | + for store_line in self._deferred_reduction_stores: |
| 2294 | + self.body.writeline(store_line) | ||
| 2295 | + self._deferred_reduction_stores.clear() | ||
| 2296 | + if self.numof_reduction_axis() > 1 and range_val.prefix == "r": | ||
| 2297 | + self.body.splice(self.stores) | ||
| 2298 | + self.stores.clear() | ||
| 2299 | + self.post_loop_combine.clear() | ||
| 2300 | + self.post_loop_store.clear() | ||
| 2273 | 2301 | ||
| 2274 | # tiling axis and but not last tiling | 2302 | # tiling axis and but not last tiling |
| 2275 | elif range_val.is_tiling_axis: | 2303 | elif range_val.is_tiling_axis: |
| @@ -2280,10 +2308,24 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2280 | indexing_code = None | 2308 | indexing_code = None |
| 2281 | if not range_val.is_no_loop_axis: | 2309 | if not range_val.is_no_loop_axis: |
| 2282 | do_indent = True | 2310 | do_indent = True |
| 2311 | + if is_first_reduction_tiling or self.numof_reduction_axis() <= 1: | ||
| 2312 | + self.body.splice(self.prefix) | ||
| 2283 | self.body.writeline( | 2313 | self.body.writeline( |
| 2284 | f"for loop_{range_val.name} in range(loops_{range_val.name}):" | 2314 | f"for loop_{range_val.name} in range(loops_{range_val.name}):" |
| 2285 | ) | 2315 | ) |
| 2286 | loop_body(index, indexing_code, is_last_axis, do_indent=do_indent) | 2316 | loop_body(index, indexing_code, is_last_axis, do_indent=do_indent) |
| 2317 | + if is_first_reduction_tiling and use_outer_reduction_post_loop: | ||
| 2318 | + self.body.splice(self.post_loop_combine) | ||
| 2319 | + for store_line in self.post_loop_store._lines: | ||
| 2320 | + self.body.writeline(store_line) | ||
| 2321 | + for store_line in self.stores._lines: | ||
| 2322 | + self.body.writeline(store_line) | ||
| 2323 | + for store_line in self._deferred_reduction_stores: | ||
| 2324 | + self.body.writeline(store_line) | ||
| 2325 | + self._deferred_reduction_stores.clear() | ||
| 2326 | + self.stores.clear() | ||
| 2327 | + self.post_loop_combine.clear() | ||
| 2328 | + self.post_loop_store.clear() | ||
| 2287 | 2329 | ||
| 2288 | elif not is_last_axis: | 2330 | elif not is_last_axis: |
| 2289 | do_indent = True | 2331 | do_indent = True |
| @@ -2324,11 +2366,20 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2324 | codegen_range(0) | 2366 | codegen_range(0) |
| 2325 | else: | 2367 | else: |
| 2326 | last_axis_order = self.tiling_axis[-1].sorted_order | 2368 | last_axis_order = self.tiling_axis[-1].sorted_order |
| 2327 | - if self.persistent_reduction and self.numof_reduction_axis() > 1: | 2369 | + skip_reduction_axes = False |
| 2370 | + if self.numof_reduction_axis() > 1: | ||
| 2328 | last_axis_order = last_axis_order - self.numof_reduction_axis() + 1 | 2371 | last_axis_order = last_axis_order - self.numof_reduction_axis() + 1 |
| 2372 | + skip_reduction_axes = not any( | ||
| 2373 | + self.find_axis_in_load_store(axis) | ||
| 2374 | + for axis in self.sorted_axis[last_axis_order:] | ||
| 2375 | + if axis.prefix == "r" | ||
| 2376 | + ) | ||
| 2329 | for _ in range(last_axis_order): | 2377 | for _ in range(last_axis_order): |
| 2330 | self.body.do_indent() | 2378 | self.body.do_indent() |
| 2331 | - codegen_range(last_axis_order) | 2379 | + if skip_reduction_axes: |
| 2380 | + write_pointwise(allow_stores=True) | ||
| 2381 | + else: | ||
| 2382 | + codegen_range(last_axis_order) | ||
| 2332 | for _ in range(last_axis_order): | 2383 | for _ in range(last_axis_order): |
| 2333 | self.body.do_unindent() | 2384 | self.body.do_unindent() |
| 2334 | 2385 | ||
| @@ -3193,6 +3244,19 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 3193 | ndims = self.triton_tensor_ndim() | 3244 | ndims = self.triton_tensor_ndim() |
| 3194 | if ndims == 1: | 3245 | if ndims == 1: |
| 3195 | return f"triton_helpers.promote_to_tensor({value})" | 3246 | return f"triton_helpers.promote_to_tensor({value})" |
| 3247 | + | ||
| 3248 | + if self.numof_reduction_axis() > 1 and self.is_contiguous_reduction(): | ||
| 3249 | + if not self.golden_var_list: | ||
| 3250 | + self.select_golden_varlist() | ||
| 3251 | + | ||
| 3252 | + dense_list = self.reduce_analysis.dense_size_list() | ||
| 3253 | + for i, axis in enumerate(reversed(self.golden_var_list)): | ||
| 3254 | + if axis.name[0] == "r": | ||
| 3255 | + dense_list[i] = "1" | ||
| 3256 | + | ||
| 3257 | + expand_str = ", ".join(dense_list) | ||
| 3258 | + return f"{value}.reshape({expand_str})" | ||
| 3259 | + | ||
| 3196 | dense_list = self.dense_size_list() | 3260 | dense_list = self.dense_size_list() |
| 3197 | dense_list[dim] = "1" | 3261 | dense_list[dim] = "1" |
| 3198 | contiguous_reduction = self.is_contiguous_reduction() | 3262 | contiguous_reduction = self.is_contiguous_reduction() |
| @@ -3317,6 +3381,8 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 3317 | self.reduce_analysis = ReductionAnalysis(self) | 3381 | self.reduce_analysis = ReductionAnalysis(self) |
| 3318 | 3382 | ||
| 3319 | dense_size_str = self.dense_size_str() | 3383 | dense_size_str = self.dense_size_str() |
| 3384 | + permute_order = None | ||
| 3385 | + need_permute = False | ||
| 3320 | axis_list = [] | 3386 | axis_list = [] |
| 3321 | for index in self.load_store_indexing: | 3387 | for index in self.load_store_indexing: |
| 3322 | for axis in V.kernel.range_tree_nodes: | 3388 | for axis in V.kernel.range_tree_nodes: |
| @@ -3332,17 +3398,39 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 3332 | ), | 3398 | ), |
| 3333 | value, | 3399 | value, |
| 3334 | ) | 3400 | ) |
| 3335 | - if len(dense_size_str) > 2 and ( | 3401 | + if ( |
| 3336 | - not self.persistent_reduction or self.numof_reduction_axis() != 1 | 3402 | + len(dense_size_str) > 2 |
| 3337 | - ): | 3403 | + and ( |
| 3338 | - value = self._map_tuple_or_scalar( | 3404 | + not self.persistent_reduction or self.numof_reduction_axis() != 1 |
| 3339 | - lambda v: self.cse.generate( | ||
| 3340 | - self.compute, | ||
| 3341 | - f"tl.reshape({v}, {dense_size_str})", | ||
| 3342 | - dtype=v.dtype, | ||
| 3343 | - ), | ||
| 3344 | - value, | ||
| 3345 | ) | 3405 | ) |
| 3406 | + ): | ||
| 3407 | + if self.numof_reduction_axis() > 1 and self.is_contiguous_reduction(): | ||
| 3408 | + value_order = list(reversed(self.golden_var_list)) | ||
| 3409 | + target_order = [x for x in value_order if x.name[0] != "r"] + [ | ||
| 3410 | + x for x in value_order if x.name[0] == "r" | ||
| 3411 | + ] | ||
| 3412 | + permute_order = [value_order.index(x) for x in target_order] | ||
| 3413 | + current_order = list(range(len(value_order))) | ||
| 3414 | + need_permute = permute_order != current_order | ||
| 3415 | + | ||
| 3416 | + if need_permute: | ||
| 3417 | + value = self._map_tuple_or_scalar( | ||
| 3418 | + lambda v: self.cse.generate( | ||
| 3419 | + self.compute, | ||
| 3420 | + f"tl.reshape({v}.permute({permute_order}), {dense_size_str})", | ||
| 3421 | + dtype=v.dtype, | ||
| 3422 | + ), | ||
| 3423 | + value, | ||
| 3424 | + ) | ||
| 3425 | + else: | ||
| 3426 | + value = self._map_tuple_or_scalar( | ||
| 3427 | + lambda v: self.cse.generate( | ||
| 3428 | + self.compute, | ||
| 3429 | + f"tl.reshape({v}, {dense_size_str})", | ||
| 3430 | + dtype=v.dtype, | ||
| 3431 | + ), | ||
| 3432 | + value, | ||
| 3433 | + ) | ||
| 3346 | 3434 | ||
| 3347 | dim: int | 3435 | dim: int |
| 3348 | root_op: str | 3436 | root_op: str |
| @@ -3382,7 +3470,10 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 3382 | torch_acc_type = upcast_acc_dtype(src_dtype) | 3470 | torch_acc_type = upcast_acc_dtype(src_dtype) |
| 3383 | result_var: Any = self.cse.newvar(dtype=torch_acc_type) | 3471 | result_var: Any = self.cse.newvar(dtype=torch_acc_type) |
| 3384 | result_var.mask_vars = {var for var in masks if var[0] != "r"} # noqa: set_linter | 3472 | result_var.mask_vars = {var for var in masks if var[0] != "r"} # noqa: set_linter |
| 3385 | - cond = f"({' & '.join(masks)}).reshape({dense_size_str})" | 3473 | + cond_expr = f"({' & '.join(masks)})" |
| 3474 | + if need_permute: | ||
| 3475 | + cond_expr = f"{cond_expr}.permute({permute_order})" | ||
| 3476 | + cond = f"{cond_expr}.reshape({dense_size_str})" | ||
| 3386 | 3477 | ||
| 3387 | def where_cond(tval, fval): | 3478 | def where_cond(tval, fval): |
| 3388 | if not cond: | 3479 | if not cond: |
请对该场景新增 ut 来看护