diff --git a/paddle/phi/kernels/impl/matmul_kernel_impl.h b/paddle/phi/kernels/impl/matmul_kernel_impl.h index 8be1016f109e7..942abbbc2db27 100644 --- a/paddle/phi/kernels/impl/matmul_kernel_impl.h +++ b/paddle/phi/kernels/impl/matmul_kernel_impl.h @@ -21,6 +21,9 @@ limitations under the License. */ #include "paddle/phi/kernels/autotune/cache_base.h" #include "paddle/phi/kernels/cast_kernel.h" #include "paddle/phi/kernels/funcs/blas/blas.h" +#include "paddle/phi/common/scalar.h" +#include "paddle/phi/core/ddim.h" +#include "paddle/phi/kernels/funcs/broadcast_function.h" #ifdef PADDLE_WITH_HIP #include "paddle/phi/kernels/funcs/blas/blaslt_impl.hip.h" #else @@ -39,6 +42,7 @@ limitations under the License. */ #include "paddle/phi/kernels/autotune/auto_tune_base.h" #endif + COMMON_DECLARE_bool(cuda_core_int8_gemm); namespace phi { @@ -2001,25 +2005,36 @@ void MatmulKernel(const Context& ctx, bool transpose_y, DenseTensor* out) { if (x.numel() == 0 || y.numel() == 0) { - auto x_dims = x.dims(); - auto y_dims = y.dims(); - if (transpose_x) { - std::swap(x_dims[x_dims.size() - 1], x_dims[x_dims.size() - 2]); - } - if (transpose_y) { - std::swap(y_dims[y_dims.size() - 1], y_dims[y_dims.size() - 2]); + std::vector x_dim_vec = common::vectorize(x.dims()); + std::vector y_dim_vec = common::vectorize(y.dims()); + + if (transpose_x && x_dim_vec.size() >= 2) { + std::swap(x_dim_vec[x_dim_vec.size() - 1], x_dim_vec[x_dim_vec.size() - 2]); } - std::vector out_dims(x_dims.size() - 1 + y_dims.size() - 1); - for (int64_t i = 0; i < x_dims.size() - 1; ++i) { - out_dims[i] = x_dims[i]; + if (transpose_y && y_dim_vec.size() >= 2) { + std::swap(y_dim_vec[y_dim_vec.size() - 1], y_dim_vec[y_dim_vec.size() - 2]); } - for (int64_t i = 1; i < y_dims.size(); ++i) { - out_dims[x_dims.size() - 1 + i - 1] = y_dims[i]; + + std::vector x_batch_dims(x_dim_vec.begin(), x_dim_vec.end() - 2); + std::vector y_batch_dims(y_dim_vec.begin(), y_dim_vec.end() - 2); + + std::vector bcast_dims; + if (!phi::funcs::BroadcastTwoVec(x_batch_dims, y_batch_dims, &bcast_dims)) { + PADDLE_THROW(phi::errors::InvalidArgument( + "Failed to broadcast input batch dimensions.")); } - out->Resize(phi::make_ddim(out_dims)); + + std::vector out_shape(bcast_dims); + int64_t m = transpose_x ? x_dim_vec.back() : x_dim_vec[x_dim_vec.size() - 2]; + int64_t n = transpose_y ? y_dim_vec[y_dim_vec.size() - 2] : y_dim_vec.back(); + out_shape.push_back(m); + out_shape.push_back(n); + + out->Resize(common::make_ddim(out_shape)); ctx.template Alloc(out); return; } + PADDLE_ENFORCE_GE( common::product(x.dims()), 0, @@ -2030,8 +2045,10 @@ void MatmulKernel(const Context& ctx, 0, common::errors::InvalidArgument( "The dims of Input(Y) should be greater than or equal to 0.")); - const std::vector x_dims = common::vectorize(x.dims()); - const std::vector y_dims = common::vectorize(y.dims()); + + const std::vector x_dims = common::vectorize(x.dims()); + const std::vector y_dims = common::vectorize(y.dims()); + MatmulJudgeDtypeKernel( ctx, x, y, x_dims, y_dims, out, transpose_x, transpose_y); } @@ -2064,7 +2081,7 @@ void MatmulWithFlattenKernelImpl(const Context& dev_ctx, #ifdef PADDLE_WITH_CUDA -template +template void MatmulWithFlattenKernelInt8Impl(const Context& dev_ctx, const DenseTensor& x, const DenseTensor& y, diff --git a/test/legacy_test/test_matmul_v3_op.py b/test/legacy_test/test_matmul_v3_op.py new file mode 100644 index 0000000000000..0a53a87c34644 --- /dev/null +++ b/test/legacy_test/test_matmul_v3_op.py @@ -0,0 +1,48 @@ +import paddle +import numpy as np + +# 所有测试用例 +test_cases = [ + # 格式: (x_shape, y_shape, expected_out_shape, dtype) + ((0, 100, 1), (0, 1, 40), (0, 100, 40), "float64"), + ((0, 100, 1), (0, 1, 4), (0, 100, 4), "float64"), + ((0, 100, 1), (1, 1, 40), (0, 100, 40), "float64"), + ((0, 100, 1), (1, 1, 4), (0, 100, 4), "float64"), + ((0, 12, 197, 197), (0, 12, 197, 64), (0, 12, 197, 64), "float16"), + ((0, 12, 197, 197), (0, 12, 197, 64), (0, 12, 197, 64), "float32"), + ((1, 0, 1), (1, 1, 40), (1, 0, 40), "float64"), + ((1, 0, 1), (1, 1, 4), (1, 0, 4), "float64"), + ((1, 100, 1), (0, 1, 40), (0, 100, 40), "float64"), + ((1, 100, 1), (0, 1, 4), (0, 100, 4), "float64"), + ((1, 100, 1), (1, 1, 0), (1, 100, 0), "float64"), + ((112, 0, 197, 197), (112, 0, 197, 64), (112, 0, 197, 64), "float16"), + ((112, 0, 197, 197), (112, 0, 197, 64), (112, 0, 197, 64), "float32"), + ((112, 12, 0, 197), (112, 12, 197, 64), (112, 12, 0, 64), "float16"), + ((112, 12, 0, 197), (112, 12, 197, 64), (112, 12, 0, 64), "float32"), + ((112, 12, 197, 197), (112, 12, 197, 0), (112, 12, 197, 0), "float16"), + ((112, 12, 197, 197), (112, 12, 197, 0), (112, 12, 197, 0), "float32"), +] + +def run_all_matmul_tests(): + for idx, (x_shape, y_shape, expected_shape, dtype) in enumerate(test_cases): + print(f"\nTest {idx+1} begin: paddle.Tensor.matmul(Tensor({x_shape}), Tensor({y_shape}), dtype={dtype})") + try: + x = paddle.zeros(x_shape, dtype=dtype) + y = paddle.zeros(y_shape, dtype=dtype) + result = x.matmul(y) + + if result.shape != expected_shape: + raise AssertionError( + f"[accuracy error] paddle.Tensor.matmul(Tensor({x_shape}), Tensor({y_shape}), dtype={dtype})\n\n" + f"Not equal to tolerance rtol=0.01, atol=0.01\n\n" + f"(shapes {result.shape}, {expected_shape} mismatch)\n" + f" x: array([], shape={x.shape}, dtype={dtype})\n" + f" y: array([], shape={y.shape}, dtype={dtype})" + ) + else: + print(f"[PASS] Shape is correct: {result.shape}") + except Exception as e: + print(f"{e}") + +if __name__ == "__main__": + run_all_matmul_tests()