已合并
[A5] [inductor] codegen error fix; bf16 fp16 layernorm lowering #26570
weizhan4创建于 2025年11月14日
[A5] [inductor] codegen error fix; bf16 fp16 layernorm lowering #26570
已合并
共 12 个文件变更+230-13
| @@ -0,0 +1,33 @@ | |||
| 1 | +import os | ||
| 2 | +import torch | ||
| 3 | +from torch import tensor, device | ||
| 4 | +from torch._dynamo.testing import rand_strided | ||
| 5 | + | ||
| 6 | + | ||
| 7 | +class Repro(torch.nn.Module): | ||
| 8 | + def __init__(self) -> None: | ||
| 9 | + super().__init__() | ||
| 10 | + | ||
| 11 | + def forward(self, embedding_list, mm, mm_3, arg54_1, arg59_1): | ||
| 12 | + cat = torch.ops.aten.cat.default(embedding_list, -1) | ||
| 13 | + sum_1 = torch.ops.aten.sum.dim_IntList(cat, [-1]) | ||
| 14 | + add_1 = torch.ops.aten.add.Tensor(sum_1, mm) | ||
| 15 | + add_3 = torch.ops.aten.add.Tensor(add_1, mm_3) | ||
| 16 | + add_4 = torch.ops.aten.add.Tensor(add_3, arg59_1) | ||
| 17 | + sigmoid = torch.ops.aten.sigmoid.default(add_4) | ||
| 18 | + return sigmoid | ||
| 19 | + | ||
| 20 | +mod = Repro().npu() | ||
| 21 | +mod = torch.compile(mod, backend="inductor", dynamic=False) | ||
| 22 | + | ||
| 23 | +if __name__ == '__main__': | ||
| 24 | + from torch._dynamo.repro.after_aot import run_repro | ||
| 25 | + with torch.no_grad(): | ||
| 26 | + arg54_1 = rand_strided((13, 1), (1, 1), device='npu', dtype=torch.float32) | ||
| 27 | + arg59_1 = rand_strided((1, 1), (1, 1), device='npu', dtype=torch.float32) | ||
| 28 | + mm = rand_strided((1, 1), (1, 1), device='npu', dtype=torch.float32) | ||
| 29 | + mm_3 = rand_strided((1, 1), (1, 1), device='npu', dtype=torch.float32) | ||
| 30 | + embedding_list = [] | ||
| 31 | + for _ in range(26, 28): | ||
| 32 | + embedding_list.append(rand_strided((1, 1, 1), (1, 1, 1), device='npu', dtype=torch.float32)) | ||
| 33 | + mod(embedding_list, mm, mm_3, arg54_1, arg59_1) | ||
| @@ -0,0 +1,23 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch import device | ||
| 3 | +import torch_npu | ||
| 4 | + | ||
| 5 | + | ||
| 6 | +class Repro(torch.nn.Module): | ||
| 7 | + def __init__(self) -> None: | ||
| 8 | + super().__init__() | ||
| 9 | + | ||
| 10 | + def forward(self): | ||
| 11 | + iota = torch.ops.prims.iota.default(50, start=0, step=1, dtype=torch.int64, device=device(type='cpu'), requires_grad=False) | ||
| 12 | + unsqueeze_23 = torch.ops.aten.unsqueeze.default(iota, 0) | ||
| 13 | + expand = torch.ops.aten.expand.default(unsqueeze_23, [16, -1]) | ||
| 14 | + device_put = torch.ops.prims.device_put.default(expand, device(type='npu', index=0)) | ||
| 15 | + convert_element_type = torch.ops.prims.convert_element_type.default(device_put, torch.int64) | ||
| 16 | + return convert_element_type | ||
| 17 | + | ||
| 18 | +mod = Repro() | ||
| 19 | +mod = torch.compile(mod, backend="inductor", dynamic=False) | ||
| 20 | + | ||
| 21 | +if __name__ == '__main__': | ||
| 22 | + with torch.no_grad(): | ||
| 23 | + mod() | ||
| @@ -0,0 +1,32 @@ | |||
| 1 | +import torch | ||
| 2 | +import torch_npu | ||
| 3 | + | ||
| 4 | + | ||
| 5 | +def main(squeeze): | ||
| 6 | + clone = torch.ops.aten.clone.default(squeeze, memory_format=torch.contiguous_format) | ||
| 7 | + select = torch.ops.aten.select.int(clone, 0, 0) | ||
| 8 | + select_1 = torch.ops.aten.select.int(clone, 0, 1) | ||
| 9 | + select_2 = torch.ops.aten.select.int(clone, 0, 2) | ||
| 10 | + view_6 = torch.ops.aten.view.default(select, [128, 2400, 16]) | ||
| 11 | + permute_2 = torch.ops.aten.permute.default(view_6, [1, 0, 2]) | ||
| 12 | + view_7 = torch.ops.aten.view.default(select_1, [128, 2400, 16]) | ||
| 13 | + permute_3 = torch.ops.aten.permute.default(view_7, [1, 0, 2]) | ||
| 14 | + view_8 = torch.ops.aten.view.default(select_2, [128, 2400, 16]) | ||
| 15 | + permute_4 = torch.ops.aten.permute.default(view_8, [1, 0, 2]) | ||
| 16 | + mul_1 = torch.ops.aten.mul.Tensor(permute_2, 0.25) | ||
| 17 | + unsqueeze_default_21 = torch.ops.aten.unsqueeze.default(mul_1, 0) | ||
| 18 | + unsqueeze_default_22 = torch.ops.aten.unsqueeze.default(permute_3, 0) | ||
| 19 | + unsqueeze_default_23 = torch.ops.aten.unsqueeze.default(permute_4, 0) | ||
| 20 | + return (unsqueeze_default_21, unsqueeze_default_22, unsqueeze_default_23) | ||
| 21 | + | ||
| 22 | +if __name__ == "__main__": | ||
| 23 | + squeeze = torch.randn((3, 128, 300, 128), device='npu', dtype=torch.float32) | ||
| 24 | + view_6 = torch.randn((128, 2400, 16), device='npu', dtype=torch.float32) | ||
| 25 | + view_7 = torch.randn((128, 2400, 16), device='npu', dtype=torch.float32) | ||
| 26 | + view_8 = torch.randn((128, 2400, 16), device='npu', dtype=torch.float32) | ||
| 27 | + arg5_1 = torch.randn((128, 128), device='npu', dtype=torch.float32) | ||
| 28 | + arg6_1 = torch.randn((128), device='npu', dtype=torch.float32) | ||
| 29 | + arg8_1 = torch.randn((384, 128), device='npu', dtype=torch.float32) | ||
| 30 | + arg7_1 = torch.randn((384), device='npu', dtype=torch.float32) | ||
| 31 | + func = torch.compile(main, backend='inductor', dynamic=False) | ||
| 32 | + func(squeeze) | ||
| @@ -1,4 +1,5 @@ | |||
| 1 | import functools | 1 | import functools |
| 2 | +from unittest import skip | ||
| 2 | import sympy | 3 | import sympy |
| 3 | import torch | 4 | import torch |
| 4 | import torch.nn.functional as F | 5 | import torch.nn.functional as F |
| @@ -10,6 +11,7 @@ from torch_npu.testing.testcase import TestCase, run_tests | |||
| 10 | 11 | ||
| 11 | 12 | ||
| 12 | class TestNpuFusionAttentionGraph(TestCase): | 13 | class TestNpuFusionAttentionGraph(TestCase): |
| 14 | + | ||
| 13 | def test_npu_graph_attention_function(self): | 15 | def test_npu_graph_attention_function(self): |
| 14 | query = torch.randn(2, 4, 8, 16, device='npu', requires_grad=True) | 16 | query = torch.randn(2, 4, 8, 16, device='npu', requires_grad=True) |
| 15 | key = torch.randn(2, 4, 8, 16, device='npu') | 17 | key = torch.randn(2, 4, 8, 16, device='npu') |
| @@ -1,4 +1,4 @@ | |||
| 1 | -from typing import List, Tuple, Dict, Any, Optional | 1 | +from typing import List, Tuple, Dict, Any, Optional, cast |
| 2 | import os | 2 | import os |
| 3 | import itertools | 3 | import itertools |
| 4 | import sympy | 4 | import sympy |
| @@ -189,8 +189,68 @@ def generate_body_indexing(body, indices): | |||
| 189 | body.generate_indirect_replacements() | 189 | body.generate_indirect_replacements() |
| 190 | 190 | ||
| 191 | 191 | ||
| 192 | +def remove_zero_terms_impl(expr, var_ranges): | ||
| 193 | + shape_env = V.graph.sizevars.shape_env | ||
| 194 | + var_to_range = dict(shape_env.var_to_range) | ||
| 195 | + var_to_range.update( | ||
| 196 | + { | ||
| 197 | + k: ValueRanges( | ||
| 198 | + 0, max(0, v - 1) if not has_free_symbols([v]) else IntInfinity() | ||
| 199 | + ) | ||
| 200 | + for k, v in var_ranges.items() | ||
| 201 | + } | ||
| 202 | + ) | ||
| 203 | + for var in expr.free_symbols: | ||
| 204 | + if var not in var_to_range: | ||
| 205 | + var_to_range[var] = ValueRanges(0, IntInfinity()) | ||
| 206 | + | ||
| 207 | + var_to_range_tuple = cast( | ||
| 208 | + tuple[tuple[sympy.Symbol, ValueRanges[sympy.Expr]]], | ||
| 209 | + tuple(var_to_range.items()), | ||
| 210 | + ) | ||
| 211 | + | ||
| 212 | + axioms = [] | ||
| 213 | + for var, upper_bound in var_ranges.items(): | ||
| 214 | + axioms.append(0 <= var) | ||
| 215 | + axioms.append(var < upper_bound) | ||
| 216 | + axioms = tuple(axioms) + shape_env.get_axioms() | ||
| 217 | + | ||
| 218 | + def statically_known(expr): | ||
| 219 | + evaluated = shape_env._maybe_evaluate_static( | ||
| 220 | + expr, | ||
| 221 | + axioms=axioms, | ||
| 222 | + var_to_range=var_to_range_tuple, | ||
| 223 | + ) | ||
| 224 | + return bool(evaluated) | ||
| 225 | + | ||
| 226 | + def _remove_zero_terms(base, divisor): | ||
| 227 | + if statically_known(base < divisor): | ||
| 228 | + return sympy.Integer(0) | ||
| 229 | + return FloorDiv(base, divisor) | ||
| 230 | + | ||
| 231 | + replacements = {} | ||
| 232 | + for sub_expr in expr.atoms(FloorDiv): | ||
| 233 | + base, divisor = sub_expr.args | ||
| 234 | + if statically_known(base < divisor): | ||
| 235 | + replacements[sub_expr] = sympy.Integer(0) | ||
| 236 | + | ||
| 237 | + if replacements: | ||
| 238 | + expr = expr.xreplace(replacements) | ||
| 239 | + | ||
| 240 | + return expr | ||
| 241 | + | ||
| 242 | + | ||
| 243 | +# Eliminate terms such as 2560(((320p1 + p2)//2560)) when (320*p1 + p2)//2560 is constantly zero | ||
| 244 | +def remove_zero_terms(indexing, var_ranges): | ||
| 245 | + for key, expr in indexing.items(): | ||
| 246 | + if expr.has(FloorDiv): | ||
| 247 | + new_expr = remove_zero_terms_impl(expr, var_ranges) | ||
| 248 | + indexing[key] = new_expr | ||
| 249 | + | ||
| 250 | + | ||
| 192 | def transform_dims_in_indexing(self, indices): | 251 | def transform_dims_in_indexing(self, indices): |
| 193 | if self.indexing is None: | 252 | if self.indexing is None: |
| 253 | + remove_zero_terms(self.indexing_exprs, self.var_ranges) | ||
| 194 | generate_body_indexing(self, indices) | 254 | generate_body_indexing(self, indices) |
| 195 | 255 | ||
| 196 | if V.kernel is not None and isinstance(V.kernel, NPUIndexTritonKernel): | 256 | if V.kernel is not None and isinstance(V.kernel, NPUIndexTritonKernel): |
| @@ -287,7 +287,6 @@ class SplitTiling: | |||
| 287 | 287 | ||
| 288 | if self.kernel.inside_reduction: | 288 | if self.kernel.inside_reduction: |
| 289 | construct_low_dim() | 289 | construct_low_dim() |
| 290 | - return | ||
| 291 | 290 | ||
| 292 | # for non-reduction, write index should be considered | 291 | # for non-reduction, write index should be considered |
| 293 | for node in self.kernel.node_schedule: | 292 | for node in self.kernel.node_schedule: |
| @@ -865,6 +865,8 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 865 | if line.find('tl.load') >= 0 and self.is_isolated_symbol(line, range_val): | 865 | if line.find('tl.load') >= 0 and self.is_isolated_symbol(line, range_val): |
| 866 | return True | 866 | return True |
| 867 | for line in self.post_loop_store._lines: | 867 | for line in self.post_loop_store._lines: |
| 868 | + if isinstance(line, DeferredLine): | ||
| 869 | + line = line.line | ||
| 868 | if line.find('tl.store') >= 0 and self.is_isolated_symbol(line, range_val): | 870 | if line.find('tl.store') >= 0 and self.is_isolated_symbol(line, range_val): |
| 869 | return True | 871 | return True |
| 870 | for line in self.stores._lines: | 872 | for line in self.stores._lines: |
| @@ -4,6 +4,7 @@ from typing import Any, Callable, Dict, Optional, TYPE_CHECKING | |||
| 4 | import torch | 4 | import torch |
| 5 | from torch._inductor import config | 5 | from torch._inductor import config |
| 6 | from triton.runtime.driver import driver | 6 | from triton.runtime.driver import driver |
| 7 | +from torch_npu.npu._backends import get_soc_version | ||
| 7 | 8 | ||
| 8 | enable_npu_indexing = True | 9 | enable_npu_indexing = True |
| 9 | 10 | ||
| @@ -78,7 +79,7 @@ acc_comp_tol = { | |||
| 78 | "default": {'rtol': 1.3e-6, 'atol': 1e-5}, | 79 | "default": {'rtol': 1.3e-6, 'atol': 1e-5}, |
| 79 | } | 80 | } |
| 80 | 81 | ||
| 81 | -if ("Ascend910B" in target.arch): | 82 | +if 220 <= get_soc_version() < 240 or get_soc_version() >= 250: |
| 82 | num_vector_core = num_cube_core * 2 | 83 | num_vector_core = num_cube_core * 2 |
| 83 | 84 | ||
| 84 | arch_support_simt = False | 85 | arch_support_simt = False |
| @@ -15,7 +15,8 @@ DECOMPOSITION_OVERLOAD_OP = [ | |||
| 15 | aten._log_softmax_backward_data, | 15 | aten._log_softmax_backward_data, |
| 16 | aten.embedding_dense_backward, | 16 | aten.embedding_dense_backward, |
| 17 | aten.addmm, | 17 | aten.addmm, |
| 18 | - aten.gelu | 18 | + aten.gelu, |
| 19 | + aten.native_layer_norm | ||
| 19 | ] | 20 | ] |
| 20 | 21 | ||
| 21 | 22 | ||
| @@ -40,14 +40,17 @@ from torch._inductor.lowering import ( | |||
| 40 | rsqrt, | 40 | rsqrt, |
| 41 | mul | 41 | mul |
| 42 | ) | 42 | ) |
| 43 | + | ||
| 43 | from torch._higher_order_ops.triton_kernel_wrap import triton_kernel_wrapper_mutation | 44 | from torch._higher_order_ops.triton_kernel_wrap import triton_kernel_wrapper_mutation |
| 44 | from torch._inductor.lowering import (unsqueeze, index_put_as_masked_fill, index_put_fallback, needs_fallback_due_to_atomic_add_limitations, view, check_and_broadcast_indices, index_output_size_and_inner_fn, expand, clone, new_empty, scatter_fallback, full_like) | 45 | from torch._inductor.lowering import (unsqueeze, index_put_as_masked_fill, index_put_fallback, needs_fallback_due_to_atomic_add_limitations, view, check_and_broadcast_indices, index_output_size_and_inner_fn, expand, clone, new_empty, scatter_fallback, full_like) |
| 45 | from torch._inductor.virtualized import V, ops | 46 | from torch._inductor.virtualized import V, ops |
| 46 | - | 47 | +from torch_npu.npu._backends import get_soc_version |
| 47 | from torch_npu import npu_dtype_cast, _npu_dtype_cast | 48 | from torch_npu import npu_dtype_cast, _npu_dtype_cast |
| 49 | +from torch_npu.npu._backends import get_soc_version | ||
| 50 | +from torch_npu._inductor import ir as npu_ir | ||
| 48 | from .ir import IndexputTemplate, ScatterTemplate | 51 | from .ir import IndexputTemplate, ScatterTemplate |
| 49 | from .lowering_op_list import GENERATE_LIST, GENERATE_LIST2, FALLBACK_LIST, LOWERING_OVERLOAD_OP | 52 | from .lowering_op_list import GENERATE_LIST, GENERATE_LIST2, FALLBACK_LIST, LOWERING_OVERLOAD_OP |
| 50 | -from .config import inductor_indirect_memory_simt_template | 53 | +from .config import inductor_indirect_memory_simt_template, lowering_cat_with_concat_kernel |
| 51 | 54 | ||
| 52 | 55 | ||
| 53 | def npu_make_fallback(op, layout_constraint=None, warn=True, override_decomp=False): | 56 | def npu_make_fallback(op, layout_constraint=None, warn=True, override_decomp=False): |
| @@ -110,6 +113,7 @@ lowering.make_reduction = make_reduction | |||
| 110 | aten = torch.ops.aten | 113 | aten = torch.ops.aten |
| 111 | tr_c10d = torch.ops.tr_c10d | 114 | tr_c10d = torch.ops.tr_c10d |
| 112 | prims = torch.ops.prims | 115 | prims = torch.ops.prims |
| 116 | +npu = torch.ops.npu | ||
| 113 | 117 | ||
| 114 | 118 | ||
| 115 | def _init_set(input_list, output_set): | 119 | def _init_set(input_list, output_set): |
| @@ -206,11 +210,11 @@ def _register_npu_inductor_fallbacks(): | |||
| 206 | return to_dtype(x, dtype, copy=True) | 210 | return to_dtype(x, dtype, copy=True) |
| 207 | return fallback_cumsum(x, dim=axis, dtype=dtype) | 211 | return fallback_cumsum(x, dim=axis, dtype=dtype) |
| 208 | 212 | ||
| 209 | - @register_lowering(npu_dtype_cast, type_promotion_kind=None) | 213 | + @register_lowering(npu.npu_dtype_cast, type_promotion_kind=None) |
| 210 | def _convert_npu_type(x: TensorBox, dtype: torch.dtype): | 214 | def _convert_npu_type(x: TensorBox, dtype: torch.dtype): |
| 211 | return to_dtype(x, dtype, copy=True) | 215 | return to_dtype(x, dtype, copy=True) |
| 212 | 216 | ||
| 213 | - @register_lowering(_npu_dtype_cast, type_promotion_kind=None) | 217 | + @register_lowering(npu._npu_dtype_cast, type_promotion_kind=None) |
| 214 | def _convert__npu_type(x: TensorBox, dtype: torch.dtype): | 218 | def _convert__npu_type(x: TensorBox, dtype: torch.dtype): |
| 215 | return to_dtype(x, dtype, copy=True) | 219 | return to_dtype(x, dtype, copy=True) |
| 216 | 220 | ||
| @@ -263,8 +267,6 @@ def _register_npu_inductor_fallbacks(): | |||
| 263 | if len(inputs) == 1: | 267 | if len(inputs) == 1: |
| 264 | return clone(inputs[0]) | 268 | return clone(inputs[0]) |
| 265 | 269 | ||
| 266 | - from torch_npu._inductor import ir as npu_ir | ||
| 267 | - from torch_npu._inductor.config import lowering_cat_with_concat_kernel | ||
| 268 | if lowering_cat_with_concat_kernel: | 270 | if lowering_cat_with_concat_kernel: |
| 269 | def is_reindex_view(x) -> bool: | 271 | def is_reindex_view(x) -> bool: |
| 270 | if isinstance(x, (TensorBox, ir.StorageBox)): | 272 | if isinstance(x, (TensorBox, ir.StorageBox)): |
| @@ -738,6 +740,63 @@ def _register_npu_inductor_fallbacks(): | |||
| 738 | inputs = [to_dtype(inp, dtype) for inp in inputs] | 740 | inputs = [to_dtype(inp, dtype) for inp in inputs] |
| 739 | return TensorBox(ir.ConcatKernel.create(inputs, dim)) | 741 | return TensorBox(ir.ConcatKernel.create(inputs, dim)) |
| 740 | 742 | ||
| 743 | + | ||
| 744 | + def native_layer_norm( | ||
| 745 | + x, | ||
| 746 | + normalized_shape, | ||
| 747 | + weight=None, | ||
| 748 | + bias=None, | ||
| 749 | + eps=1e-5 | ||
| 750 | + ): | ||
| 751 | + # Performance consideration: fallback for bfloat16 and float16 | ||
| 752 | + if get_soc_version() >= 250 and \ | ||
| 753 | + (x.dtype == torch.bfloat16 or x.dtype == torch.float16): | ||
| 754 | + return fallback_handler(aten.native_layer_norm.default)(x, normalized_shape, weight, bias, eps) | ||
| 755 | + # Validate input | ||
| 756 | + if not isinstance(normalized_shape, (list, tuple)): | ||
| 757 | + normalized_shape = (normalized_shape,) | ||
| 758 | + | ||
| 759 | + normalized_ndim = len(normalized_shape) | ||
| 760 | + input_shape = x.get_size() | ||
| 761 | + | ||
| 762 | + # Calculate reduction dimension indices | ||
| 763 | + reduce_dims = list(range(len(input_shape) - normalized_ndim, len(input_shape))) | ||
| 764 | + | ||
| 765 | + # Compute mean and variance | ||
| 766 | + var, mean = var_mean_helper_( | ||
| 767 | + x=x, | ||
| 768 | + axis=reduce_dims, | ||
| 769 | + correction=0, # Layer normalization uses 0 correction (population variance) | ||
| 770 | + keepdim=True, # Keep dimensions for broadcasting | ||
| 771 | + return_mean=True | ||
| 772 | + ) | ||
| 773 | + | ||
| 774 | + # Calculate normalized result (x - mean) / sqrt(var + eps) | ||
| 775 | + x_normalized = sub(x, mean) | ||
| 776 | + | ||
| 777 | + # Add eps to variance | ||
| 778 | + eps_tensor = ir.IndexingConstant(index=eps, dtype=var.get_dtype(), device=var.get_device()) | ||
| 779 | + eps_tensor = ExpandView.create(eps_tensor, var.get_size()) | ||
| 780 | + var_eps = add(var, eps_tensor) | ||
| 781 | + | ||
| 782 | + # Calculate reciprocal of sqrt(var + eps) | ||
| 783 | + inv_std = rsqrt(var_eps) # 1 / sqrt(var + eps) | ||
| 784 | + | ||
| 785 | + # Normalization | ||
| 786 | + normalized = mul(x_normalized, inv_std) | ||
| 787 | + | ||
| 788 | + # Apply optional affine transformation (gamma * normalized + beta) | ||
| 789 | + if weight is not None: | ||
| 790 | + # weight will be broadcast automatically, mul function in lowering supports broadcasting | ||
| 791 | + normalized = mul(normalized, weight) | ||
| 792 | + | ||
| 793 | + if bias is not None: | ||
| 794 | + # add will be broadcast automatically | ||
| 795 | + normalized = add(normalized, bias) | ||
| 796 | + | ||
| 797 | + # native_layer_norm returns three values: output, mean, reciprocal of standard deviation | ||
| 798 | + return normalized, mean, inv_std | ||
| 799 | + | ||
| 741 | make_fallback(aten._log_softmax) | 800 | make_fallback(aten._log_softmax) |
| 742 | make_fallback(aten.nll_loss_forward) | 801 | make_fallback(aten.nll_loss_forward) |
| 743 | 802 | ||
| @@ -1,14 +1,16 @@ | |||
| 1 | import torch | 1 | import torch |
| 2 | from torch._higher_order_ops.triton_kernel_wrap import triton_kernel_wrapper_mutation | 2 | from torch._higher_order_ops.triton_kernel_wrap import triton_kernel_wrapper_mutation |
| 3 | -from torch_npu import npu_dtype_cast, _npu_dtype_cast | ||
| 4 | from torch_npu._inductor.config import arch_support_simt | 3 | from torch_npu._inductor.config import arch_support_simt |
| 5 | from .config import inductor_indirect_memory_simt_template, inductor_support_simt | 4 | from .config import inductor_indirect_memory_simt_template, inductor_support_simt |
| 6 | 5 | ||
| 7 | aten = torch.ops.aten | 6 | aten = torch.ops.aten |
| 8 | tr_c10d = torch.ops.tr_c10d | 7 | tr_c10d = torch.ops.tr_c10d |
| 9 | prims = torch.ops.prims | 8 | prims = torch.ops.prims |
| 9 | +npu = torch.ops.npu | ||
| 10 | 10 | ||
| 11 | GENERATE_LIST = [ | 11 | GENERATE_LIST = [ |
| 12 | + aten.copy_, | ||
| 13 | + prims.device_put, | ||
| 12 | prims.iota, | 14 | prims.iota, |
 | |||
| 13 | aten.full, | 15 | aten.full, |
| 14 | aten.mul, | 16 | aten.mul, |
| @@ -58,8 +60,8 @@ GENERATE_LIST = [ | |||
| 58 | aten.clamp, | 60 | aten.clamp, |
| 59 | aten.clamp_max, | 61 | aten.clamp_max, |
| 60 | aten.mean, | 62 | aten.mean, |
| 61 | - npu_dtype_cast, | 63 | + npu.npu_dtype_cast, |
| 62 | - _npu_dtype_cast, | 64 | + npu._npu_dtype_cast, |
| 63 | aten.select_scatter, | 65 | aten.select_scatter, |
| 64 | aten.slice_scatter, | 66 | aten.slice_scatter, |
| 65 | prims.broadcast_in_dim, | 67 | prims.broadcast_in_dim, |
| @@ -80,6 +82,7 @@ GENERATE_LIST = [ | |||
| 80 | aten.reciprocal, | 82 | aten.reciprocal, |
| 81 | aten._assert_scalar, | 83 | aten._assert_scalar, |
| 82 | triton_kernel_wrapper_mutation, | 84 | triton_kernel_wrapper_mutation, |
| 85 | + aten.native_layer_norm, | ||
| 83 | ] | 86 | ] |
| 84 | 87 | ||
| 85 | GENERATE_LIST2 = [ | 88 | GENERATE_LIST2 = [ |
| @@ -1139,6 +1139,8 @@ def triton_config_npu_index( | |||
| 1139 | cfg = {} | 1139 | cfg = {} |
| 1140 | for x in split_axis: | 1140 | for x in split_axis: |
| 1141 | cfg[f"{axis_names[x].upper()}BLOCK"] = size_hints[x] | 1141 | cfg[f"{axis_names[x].upper()}BLOCK"] = size_hints[x] |
| 1142 | + for x in tiling_axis: | ||
| 1143 | + cfg[f"{axis_names[x].upper()}BLOCK_SUB"] = size_hints[x] | ||
| 1142 | if not cfg: | 1144 | if not cfg: |
| 1143 | cfg["dummy"] = 1 | 1145 | cfg["dummy"] = 1 |
| 1144 | tmp = Config(cfg, num_warps=num_warps, num_stages=num_stages) | 1146 | tmp = Config(cfg, num_warps=num_warps, num_stages=num_stages) |
这部分修改也会影响到A3吧？