diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 365ec64b04..cfe46b4f13 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -63,7 +63,6 @@ _materialize_integer_literal, _normalize_address_space, _resolve, - _strip_integer_signedness, mask_type, part_tensor_view_type, part_tensor_view_type_from_dims, @@ -2032,13 +2031,12 @@ def _infer_vdup_scalar_result_type(input_value, mask_value, *, context: str): scalar_type = scalar_raw.type mask_bits = _mask_granularity_bits(mask_value, context=context) if IntegerType.isinstance(scalar_type): - scalar_type = _strip_integer_signedness(scalar_raw) - scalar_width = IntegerType(scalar_type.type).width + scalar_width = IntegerType(scalar_type).width if scalar_width != mask_bits: raise TypeError( f"{context} expects scalar input width {scalar_width} to match mask granularity b{mask_bits}" ) - element_type = scalar_type.type + element_type = scalar_type elif F16Type.isinstance(scalar_type) or BF16Type.isinstance(scalar_type): if mask_bits != 16: raise TypeError(f"{context} expects f16/bf16 scalar input to pair with mask_b16, got mask_b{mask_bits}") @@ -2090,6 +2088,12 @@ def vdup(input_value, mask, position=None): raise TypeError("vdup(scalar, mask, position=...) does not support position; position is only valid for vector input") raw_input = _coerce_vdup_scalar_input(input_value, mask, context="vdup(scalar, mask)") result_type = _infer_vdup_scalar_result_type(raw_input, mask, context="vdup(scalar, mask)") + result_element_type = _pto.VRegType(result_type).element_type + raw_input = coerce_scalar_to_type( + raw_input, + result_element_type, + context="vdup(scalar, mask)", + ) normalized_position = None return wrap_surface_value( _pto.VdupOp( diff --git a/ptodsl/tests/test_issue_1102_vdup.py b/ptodsl/tests/test_issue_1102_vdup.py new file mode 100644 index 0000000000..83d8b749e8 --- /dev/null +++ b/ptodsl/tests/test_issue_1102_vdup.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +from ptodsl import pto + + +def _compile_vdup(name, dtype, mask_bits, value): + @pto.jit(name=name, kernel_kind="vector", target="a5", mode="explicit") + def kernel(): + mask = getattr(pto, f"pset_b{mask_bits}")(pto.MaskPattern.ALL) + pto.vdup(pto.const(value, dtype=dtype), mask) + + return kernel.compile().mlir_text() + + +def _compile_vdup_consumer(): + @pto.jit(name="issue_1102_vdup_vor", kernel_kind="vector", target="a5", mode="explicit") + def kernel(): + base = pto.const(0, dtype=pto.ui64) + source = pto.castptr(base, pto.ptr(pto.ui16, "ub")) + mask = pto.pset_b16(pto.MaskPattern.ALL) + loaded = pto.vlds(source, pto.const(0)) + duplicated = pto.vdup(pto.ui16(1), mask) + pto.vor(loaded, duplicated, mask) + + return kernel.compile().mlir_text() + + +def main(): + cases = ( + ("i8", pto.i8, 8, 1, "i8, !pto.mask -> !pto.vreg<256xi8>"), + ("si8", pto.si8, 8, 1, "si8, !pto.mask -> !pto.vreg<256xsi8>"), + ("ui8", pto.ui8, 8, 1, "ui8, !pto.mask -> !pto.vreg<256xui8>"), + ("i16", pto.i16, 16, 1, "i16, !pto.mask -> !pto.vreg<128xi16>"), + ("si16", pto.si16, 16, 1, "si16, !pto.mask -> !pto.vreg<128xsi16>"), + ("ui16", pto.ui16, 16, 1, "ui16, !pto.mask -> !pto.vreg<128xui16>"), + ("i32", pto.i32, 32, 1, "i32, !pto.mask -> !pto.vreg<64xi32>"), + ("si32", pto.si32, 32, 1, "si32, !pto.mask -> !pto.vreg<64xsi32>"), + ("ui32", pto.ui32, 32, 1, "ui32, !pto.mask -> !pto.vreg<64xui32>"), + ) + for name, dtype, mask_bits, value, expected_type in cases: + text = _compile_vdup(f"issue_1102_vdup_{name}", dtype, mask_bits, value) + if expected_type not in text: + raise AssertionError( + f"{name} vdup should preserve the scalar integer signedness in its result vector; " + f"expected {expected_type!r} in:\n{text}" + ) + _compile_vdup_consumer() + print("issue_1102_vdup: PASS") + + +if __name__ == "__main__": + main()