已开启
fix concat convert to fp32 #3163
xgw创建于 6月4日
fix concat convert to fp32 #3163
已开启
共 1 个文件变更+8-0
| @@ -693,6 +693,8 @@ class NPUKernel(Kernel): | |||
| 693 | """ | 693 | """ |
| 694 | axis = concat_buffer.axis | 694 | axis = concat_buffer.axis |
| 695 | out_name = concat_buffer.get_name() | 695 | out_name = concat_buffer.get_name() |
| 696 | + out_dtype = concat_buffer.get_dtype() | ||
| 697 | + need_cast = out_dtype in {torch.bfloat16, torch.float16} | ||
| 696 | 698 | ||
| 697 | loaded = [] | 699 | loaded = [] |
| 698 | for buf, input_layout in zip(concat_buffer.inputs, concat_buffer.input_layouts): | 700 | for buf, input_layout in zip(concat_buffer.inputs, concat_buffer.input_layouts): |
| @@ -700,6 +702,8 @@ class NPUKernel(Kernel): | |||
| 700 | if in_name in self._local_stores: | 702 | if in_name in self._local_stores: |
| 701 | # 上游融合命中：取已经在 partial loop 上算好的张量 | 703 | # 上游融合命中：取已经在 partial loop 上算好的张量 |
| 702 | value, _src_loop = self._local_stores[in_name] | 704 | value, _src_loop = self._local_stores[in_name] |
| 705 | + if need_cast: | ||
| 706 | + value = ir.cast(value, dst=torch.float32, loop=_src_loop.copy().contiguous_()) | ||
| 703 | loaded.append(value) | 707 | loaded.append(value) |
| 704 | continue | 708 | continue |
| 705 | asc_in = ASCBuffer(in_name, input_layout) | 709 | asc_in = ASCBuffer(in_name, input_layout) |
| @@ -707,9 +711,13 @@ class NPUKernel(Kernel): | |||
| 707 | data, _ = self._load_buffer(in_name, partial_loop) | 711 | data, _ = self._load_buffer(in_name, partial_loop) |
| 708 | offset = partial_loop.zero_offset_() | 712 | offset = partial_loop.zero_offset_() |
| 709 | load = ir.load(data, offset=offset, loop=partial_loop) | 713 | load = ir.load(data, offset=offset, loop=partial_loop) |
| 714 | + if need_cast: | ||
| 715 | + load = ir.cast(load, dst=torch.float32, loop=partial_loop.copy().contiguous_()) | ||
| 710 | loaded.append(load) | 716 | loaded.append(load) |
| 711 | 717 | ||
| 712 | out = ir.concat(loaded, axis=axis) | 718 | out = ir.concat(loaded, axis=axis) |
| 719 | + if need_cast: | ||
| 720 | + out = ir.cast(out, dst=out_dtype, loop=self.contiguous_loop.copy().contiguous_()) | ||
| 713 | self._store_buffer(out_name, out, self.contiguous_loop) | 721 | self._store_buffer(out_name, out, self.contiguous_loop) |
| 714 | 722 | ||
| 715 | def _make_partial_loop(self, asc_in: 'ASCBuffer', concat_axis: int) -> Loop: | 723 | def _make_partial_loop(self, asc_in: 'ASCBuffer', concat_axis: int) -> Loop: |