diff --git a/backends/native/runtime/cpu/CPUEngineExecutable.cpp b/backends/native/runtime/cpu/CPUEngineExecutable.cpp index a24d7b63d87..0d2c9aedbfb 100644 --- a/backends/native/runtime/cpu/CPUEngineExecutable.cpp +++ b/backends/native/runtime/cpu/CPUEngineExecutable.cpp @@ -66,8 +66,154 @@ void validate_supported_method(const Method& method) { } } +[[noreturn]] void fail_matmul(const char* reason) { + throw std::runtime_error(std::string("cpu: matmul ") + reason); +} + +void lower_matmul(Graph& graph, NodeId id) { + const auto node = graph.node(id); + if (node.inputs.size() != 2 || node.outputs.size() != 1 || + node.outputs[0].kind != OutputValueKind::Tensor || + node.inputs[0].arg.kind() != ArgKind::Tensor || + node.inputs[1].arg.kind() != ArgKind::Tensor) { + fail_matmul("requires two tensor inputs and one tensor output"); + } + const auto left_id = node.inputs[0].arg.as_tensor().id; + const auto right_id = node.inputs[1].arg.as_tensor().id; + const auto& left_value = graph.value(left_id); + const auto& right_value = graph.value(right_id); + if (!left_value.is_tensor() || !right_value.is_tensor() || + left_value.tensor_meta().dtype != ScalarType::Float || + right_value.tensor_meta().dtype != ScalarType::Float) { + fail_matmul("requires FP32 operands"); + } + const auto left = left_value.tensor_meta().sizes; + const auto right = right_value.tensor_meta().sizes; + if (left.size() < 2 || right.size() < 2 || + std::ranges::any_of(left, [](int64_t size) { return size <= 0; }) || + std::ranges::any_of(right, [](int64_t size) { return size <= 0; })) { + fail_matmul("requires positive matrix shapes of rank at least two"); + } + const size_t rank = std::max(left.size(), right.size()); + std::vector expanded_left(rank, 1), expanded_right(rank, 1); + std::copy(left.begin(), left.end(), expanded_left.end() - left.size()); + std::copy(right.begin(), right.end(), expanded_right.end() - right.size()); + if (expanded_left[rank - 1] != expanded_right[rank - 2]) { + fail_matmul("has incompatible contraction dimensions"); + } + std::vector output_shape(rank); + int64_t batches = 1; + for (size_t dim = 0; dim + 2 < rank; ++dim) { + const auto a = expanded_left[dim]; + const auto b = expanded_right[dim]; + if (a != b && a != 1 && b != 1) { + fail_matmul("has incompatible broadcast dimensions"); + } + const auto extent = std::max(a, b); + if (batches > std::numeric_limits::max() / extent) { + fail_matmul("has an overflowing batch extent"); + } + batches *= extent; + output_shape[dim] = extent; + expanded_left[dim] = extent; + expanded_right[dim] = extent; + } + const auto rows = left[left.size() - 2]; + const auto inner = left.back(); + const auto columns = right.back(); + output_shape[rank - 2] = rows; + output_shape[rank - 1] = columns; + const auto& output = graph.value(node.outputs[0].value_id); + if (!output.is_tensor() || output.tensor_meta().dtype != ScalarType::Float || + output.tensor_meta().sizes != output_shape) { + fail_matmul("has incompatible output metadata"); + } + if (rank == 2 || + (rank == 3 && left.size() == 3 && right.size() == 3 && + left[0] == right[0])) { + auto& replacement = graph.node(id); + replacement.target = + rank == 2 ? "torch.ops.aten.mm.default" : "torch.ops.aten.bmm.default"; + replacement.inputs[1].name = "mat2"; + return; + } + const auto tensor = [](ValueId value) -> NamedArgument { + return {.name = "self", .arg = TensorArg{value}}; + }; + const auto call = [&](const char* suffix, + const char* target, + std::vector arguments, + std::vector shape) { + const std::string name = node.name + "__cpu_matmul_" + suffix; + const auto value = + graph.append_value(Value(name, ScalarType::Float, std::move(shape))); + graph.insert_node_before( + id, + Node{ + .name = name, + .op_kind = OpKind::CallFunction, + .target = target, + .inputs = std::move(arguments), + .outputs = + {{.kind = OutputValueKind::Tensor, + .value_id = value, + .elem_ids = {}}}, + .attrs = {}, + }); + return value; + }; + const auto expand = [&](const char* suffix, + ValueId source, + const std::vector& original, + const std::vector& shape) { + if (original == shape) { + return source; + } + return call( + suffix, + "torch.ops.aten.expand_copy.default", + {tensor(source), + {.name = "size", .arg = IntListArg{.values = shape, .ids = {}}}, + {.name = "implicit", .arg = BoolArg{false}}}, + shape); + }; + const auto reshape = [&](const char* suffix, + ValueId source, + const std::vector& shape) { + return call( + suffix, + "torch.ops.aten.view_copy.default", + {tensor(source), + {.name = "size", .arg = IntListArg{.values = shape, .ids = {}}}}, + shape); + }; + const auto expanded_a = expand("left_expand", left_id, left, expanded_left); + const auto expanded_b = + expand("right_expand", right_id, right, expanded_right); + const auto a = reshape("left", expanded_a, {batches, rows, inner}); + const auto b = reshape("right", expanded_b, {batches, inner, columns}); + const auto result = call( + "result", + "torch.ops.aten.bmm.default", + {tensor(a), {.name = "mat2", .arg = TensorArg{b}}}, + {batches, rows, columns}); + auto& replacement = graph.node(id); + replacement.target = "torch.ops.aten.view_copy.default"; + replacement.inputs = { + tensor(result), + {.name = "size", .arg = IntListArg{.values = output_shape, .ids = {}}}}; + graph.rebuild_def_use(); +} + void lower_native_method(Graph& graph) { functionalize_aliases(graph); + const auto schedule = graph.schedule; + for (const auto id : schedule) { + if (graph.node(id).is_call() && + graph.node(id).target == "torch.ops.aten.matmul.default") { + lower_matmul(graph, id); + } + } for (ValueId id = 0; std::cmp_less(id, graph.values.size()); ++id) { auto& value = graph.value(id); if (value.is_tensor() && value.role == ValueRole::Intermediate && @@ -175,9 +321,10 @@ class CPUEngineExecutable final : public EngineExecutable { const Package& package, CPUConstantStore& constants, cpu::RuntimeConfiguration configuration) - : method_(method), buffers_(method.graph.values.size()) { + : method_(method) { validate_supported_method(method_); lower_native_method(method_.graph); + buffers_.resize(method_.graph.values.size()); plan_ = std::make_unique( method_.graph, buffers_, diff --git a/backends/native/runtime/cpu/README.md b/backends/native/runtime/cpu/README.md index cf5de58dfb5..209a7488dae 100644 --- a/backends/native/runtime/cpu/README.md +++ b/backends/native/runtime/cpu/README.md @@ -46,6 +46,11 @@ including graph inputs and outputs, must be a contiguous tensor of rank at most Unsupported operators, including retained runtime assertions, fail provider selection. +Static FP32 `matmul` with positive operands of rank at least two is lowered to +`mm` or `bmm`, using managed copies for batch broadcasting and flattening. +Contraction dimensions, broadcast dimensions and output metadata are checked +before compilation. Vector and empty-matrix cases remain unsupported. + Direct engine clients set every input before the first execution. Inputs are written once into dedicated buffers outside the activation arena, so their values remain available across repeated executions; setting an input diff --git a/backends/native/test/runtime/cpu/test_cpu_engine.cpp b/backends/native/test/runtime/cpu/test_cpu_engine.cpp index 64fc03bbce5..896d593b39f 100644 --- a/backends/native/test/runtime/cpu/test_cpu_engine.cpp +++ b/backends/native/test/runtime/cpu/test_cpu_engine.cpp @@ -242,6 +242,21 @@ class TestGraph { &out)); } + void matmul(const char* name, Offset metadata) { + tensor(name, metadata); + const std::vector> args{ + tensor_arg("self", "x"), tensor_arg("other", "y")}; + const std::vector> out{ + fbs::CreateOutputDirect(builder, name)}; + nodes.push_back(fbs::CreateNodeDirect( + builder, + name, + fbs::OpKind::CALL_FUNCTION, + "torch.ops.aten.matmul.default", + &args, + &out)); + } + void to_copy( const char* name, const char* input, @@ -843,6 +858,157 @@ TEST(CPUEngineTest, MaterializesNativePermuteAliases) { } } +TEST(CPUEngineTest, LowersMatrixAndBatchedMatrixMultiply) { + struct Case { + size_t batches; + std::vector left; + std::vector right; + std::vector output; + }; + for (const auto& [batches, left, right, out] : + {Case{1, {2, 3}, {3, 2}, {2, 2}}, + Case{2, {2, 2, 3}, {2, 3, 2}, {2, 2, 2}}, + Case{4, {2, 2, 2, 3}, {2, 2, 3, 2}, {2, 2, 2, 2}}}) { + SCOPED_TRACE(::testing::PrintToString(out)); + TestGraph graph; + graph.input("x", graph.meta(fbs::ScalarType::FLOAT, left)); + graph.input("y", graph.meta(fbs::ScalarType::FLOAT, right)); + graph.matmul("product", graph.meta(fbs::ScalarType::FLOAT, out)); + graph.output("product"); + auto compiled = compile(graph.program()); + std::vector x(batches * 6), y(batches * 6), expected(batches * 4); + for (size_t i = 0; i < x.size(); ++i) { + x[i] = (static_cast(i % 7) - 3) * 0.25f; + y[i] = (static_cast(i % 5) - 2) * 0.5f; + } + for (size_t batch = 0; batch < batches; ++batch) { + for (size_t row = 0; row < 2; ++row) { + for (size_t column = 0; column < 2; ++column) { + float sum = 0; + for (size_t inner = 0; inner < 3; ++inner) { + sum += x[batch * 6 + row * 3 + inner] * + y[batch * 6 + inner * 2 + column]; + } + expected[batch * 4 + row * 2 + column] = sum; + } + } + } + compiled.executable->set_input(0, x.data(), x.size(), ScalarType::Float); + compiled.executable->set_input(1, y.data(), y.size(), ScalarType::Float); + for (int repeat = 0; repeat < 2; ++repeat) { + compiled.executable->execute(); + std::vector actual(expected.size()); + compiled.executable->get_output( + 0, actual.data(), actual.size(), ScalarType::Float); + EXPECT_EQ(actual, expected); + } + } +} + +TEST(CPUEngineTest, LowersMatmulWithBroadcastedBatchAndHeadDimensions) { + TestGraph graph; + graph.input("x", graph.meta(fbs::ScalarType::FLOAT, {2, 1, 3, 4})); + graph.input("y", graph.meta(fbs::ScalarType::FLOAT, {1, 5, 4, 2})); + graph.matmul("product", graph.meta(fbs::ScalarType::FLOAT, {2, 5, 3, 2})); + graph.output("product"); + auto compiled = compile(graph.program()); + std::vector x(24), y(40), expected(60); + for (size_t i = 0; i < x.size(); ++i) { + x[i] = (static_cast(i % 7) - 3) * 0.125f; + } + for (size_t i = 0; i < y.size(); ++i) { + y[i] = (static_cast(i % 5) - 2) * 0.25f; + } + for (size_t batch = 0; batch < 2; ++batch) { + for (size_t head = 0; head < 5; ++head) { + for (size_t row = 0; row < 3; ++row) { + for (size_t column = 0; column < 2; ++column) { + float sum = 0; + for (size_t inner = 0; inner < 4; ++inner) { + sum += x[batch * 12 + row * 4 + inner] * + y[head * 8 + inner * 2 + column]; + } + expected[(batch * 5 + head) * 6 + row * 2 + column] = sum; + } + } + } + } + compiled.executable->set_input(0, x.data(), x.size(), ScalarType::Float); + compiled.executable->set_input(1, y.data(), y.size(), ScalarType::Float); + for (int repeat = 0; repeat < 2; ++repeat) { + compiled.executable->execute(); + std::vector actual(expected.size()); + compiled.executable->get_output( + 0, actual.data(), actual.size(), ScalarType::Float); + EXPECT_EQ(actual, expected); + } +} + +TEST(CPUEngineTest, LowersMatmulWithPrependedBatchDimensions) { + TestGraph graph; + graph.input("x", graph.meta(fbs::ScalarType::FLOAT, {3, 4})); + graph.input("y", graph.meta(fbs::ScalarType::FLOAT, {2, 4, 2})); + graph.matmul("product", graph.meta(fbs::ScalarType::FLOAT, {2, 3, 2})); + graph.output("product"); + auto compiled = compile(graph.program()); + std::vector x(12), y(16), expected(12); + for (size_t i = 0; i < x.size(); ++i) { + x[i] = static_cast(static_cast(i % 7) - 3); + } + for (size_t i = 0; i < y.size(); ++i) { + y[i] = (static_cast(i % 5) - 2) * 0.5f; + } + for (size_t batch = 0; batch < 2; ++batch) { + for (size_t row = 0; row < 3; ++row) { + for (size_t column = 0; column < 2; ++column) { + float sum = 0; + for (size_t inner = 0; inner < 4; ++inner) { + sum += x[row * 4 + inner] * y[batch * 8 + inner * 2 + column]; + } + expected[batch * 6 + row * 2 + column] = sum; + } + } + } + compiled.executable->set_input(0, x.data(), x.size(), ScalarType::Float); + compiled.executable->set_input(1, y.data(), y.size(), ScalarType::Float); + for (int repeat = 0; repeat < 2; ++repeat) { + compiled.executable->execute(); + std::vector actual(expected.size()); + compiled.executable->get_output( + 0, actual.data(), actual.size(), ScalarType::Float); + EXPECT_EQ(actual, expected); + } +} + +TEST(CPUEngineTest, RejectsUnsupportedMatmulMetadata) { + for (const std::string scenario : + {"dtype", "rank", "inner", "batch", "output", "zero"}) { + SCOPED_TRACE(scenario); + std::vector left{1, 2, 3}, right{1, 3, 2}, out{1, 2, 2}; + if (scenario == "rank") { + left = {3}; + } else if (scenario == "inner") { + right[1] = 4; + } else if (scenario == "batch") { + left[0] = 2; + right[0] = 3; + } else if (scenario == "output") { + out[1] = 3; + } else if (scenario == "zero") { + left[2] = 0; + right[1] = 0; + } + TestGraph graph; + const auto dtype = + scenario == "dtype" ? fbs::ScalarType::LONG : fbs::ScalarType::FLOAT; + graph.input("x", graph.meta(dtype, left)); + graph.input("y", graph.meta(dtype, right)); + graph.matmul("product", graph.meta(dtype, out)); + graph.output("product"); + expect_error([&] { compile(graph.program()); }, "cpu: matmul"); + } +} + TEST(CPUEngineTest, MaterializesContiguousNativeStridedReshape) { TestGraph graph; graph.input("x", graph.meta(fbs::ScalarType::FLOAT, {2, 2}));