From 55cf3338e669ce3a89036081a7cc08c016f25117 Mon Sep 17 00:00:00 2001 From: Eric Lippert Date: Tue, 20 Sep 2022 12:16:33 -0700 Subject: [PATCH] Fix C++ code generation for typed eigen matrices. Summary: Eigen distinguishes between matrices of bools, naturals and doubles; when generating a C++ fragment to construct a BMG graph, we always generated a double matrix, but that's bad; it will not type check when compiled by the C++ compiler since the BMG methods for constructing constant matrices require an input of the appropriate type. Reviewed By: gafter, AishwaryaSivaraman Differential Revision: D39632858 fbshipit-source-id: aef37f30537fc690ed7756822a373f7411559669 --- .../ppl/compiler/fix_observations.py | 1 + src/beanmachine/ppl/compiler/gen_bmg_cpp.py | 52 ++++++++++++------- .../ppl/compiler/tests/bmg_query_test.py | 8 ++- 3 files changed, 39 insertions(+), 22 deletions(-) diff --git a/src/beanmachine/ppl/compiler/fix_observations.py b/src/beanmachine/ppl/compiler/fix_observations.py index 1b78dd05a1..5332dca717 100644 --- a/src/beanmachine/ppl/compiler/fix_observations.py +++ b/src/beanmachine/ppl/compiler/fix_observations.py @@ -52,6 +52,7 @@ def observations_fixer(bmg: BMGraphBuilder): else: # TODO: How should we deal with observations of # TODO: matrix-valued samples? + # When this gets fixed, also fix _add_observation in gen_bmg_cpp. pass # TODO: Handle the case where there are two inconsistent # TODO: observations of the same sample diff --git a/src/beanmachine/ppl/compiler/gen_bmg_cpp.py b/src/beanmachine/ppl/compiler/gen_bmg_cpp.py index 457f2d02e7..a8befc2d60 100644 --- a/src/beanmachine/ppl/compiler/gen_bmg_cpp.py +++ b/src/beanmachine/ppl/compiler/gen_bmg_cpp.py @@ -18,21 +18,29 @@ from beanmachine.ppl.compiler.lattice_typer import LatticeTyper -def _value_to_cpp_eigen(value: torch.Tensor, variable: str) -> str: +def _value_to_cpp_eigen(value: torch.Tensor, variable: str, array_type: str) -> str: # Torch tensors are row major but Eigen matrices are column major; # a torch Dirichlet distribution expects a row of parameters; # BMG expects a column. That's why we swap rows with columns here. r, c = _size_to_rc(value.size()) v = value.reshape(r, c).transpose(0, 1).contiguous() values = ", ".join(str(element) for element in v.reshape(-1).tolist()) - return f"Eigen::MatrixXd {variable}({c}, {r});\n{variable} << {values};" + return f"{array_type} {variable}({c}, {r});\n{variable} << {values};" + + +def _value_to_cpp_real_eigen(value: torch.Tensor, variable: str) -> str: + return _value_to_cpp_eigen(value, variable, "Eigen::MatrixXd") + + +def _value_to_cpp_natural_eigen(value: torch.Tensor, variable: str) -> str: + return _value_to_cpp_eigen(value, variable, "Eigen::MatrixXn") + + +def _value_to_cpp_bool_eigen(value: torch.Tensor, variable: str) -> str: + return _value_to_cpp_eigen(value, variable, "Eigen::MatrixXb") def _value_to_cpp(value: Any) -> str: - if isinstance(value, torch.Tensor): - values = ",".join(str(element) for element in value.reshape(-1).tolist()) - dims = ",".join(str(dim) for dim in value.shape) - return f"torch::from_blob((float[]){{{values}}}, {{{dims}}})" return str(value).lower() @@ -58,7 +66,10 @@ def _add_observation(self, node: bn.Observation) -> None: if isinstance(v, torch.Tensor): o = f"o{self._observation_count}" self._observation_count += 1 - s = _value_to_cpp_eigen(v, o) + # TODO: What if its not a real-valued observation? + # We do not handle this case correctly in fix_observations. + # When we do, fix this here too. + s = _value_to_cpp_real_eigen(v, o) self._code.append(s) self._code.append(f"g.observe(n{graph_id}, {o});") else: @@ -140,32 +151,33 @@ def _add_constant(self, node: bn.ConstantNode) -> None: # noqa elif t is bn.RealNode: f = f"add_constant_real({_value_to_cpp(float(v))})" elif t is bn.ConstantTensorNode: - # TODO: This should be dead code; this is a holdover - # from before BMG used Eigen as its matrix library. - # Remove it, and remove the case from _value_to_cpp. - f = f"add_constant({_value_to_cpp(v)})" + # This only happens in the rare case where we have a functional + # which returns a constant tensor and there's a query on it. + # We never rewrite that to another kind of node. + # TODO: Consider turning that into a constant of some more + # specific BMG node type. + m = _value_to_cpp_real_eigen(v, f"m{graph_id}") + f = f"add_constant_real_matrix(m{graph_id})" elif t is bn.ConstantPositiveRealMatrixNode: - m = _value_to_cpp_eigen(v, f"m{graph_id}") + m = _value_to_cpp_real_eigen(v, f"m{graph_id}") f = f"add_constant_pos_matrix(m{graph_id})" elif t is bn.ConstantRealMatrixNode: - m = _value_to_cpp_eigen(v, f"m{graph_id}") + m = _value_to_cpp_real_eigen(v, f"m{graph_id}") f = f"add_constant_real_matrix(m{graph_id})" elif t is bn.ConstantNegativeRealMatrixNode: - m = _value_to_cpp_eigen(v, f"m{graph_id}") + m = _value_to_cpp_real_eigen(v, f"m{graph_id}") f = f"add_constant_neg_matrix(m{graph_id})" elif t is bn.ConstantProbabilityMatrixNode: - m = _value_to_cpp_eigen(v, f"m{graph_id}") + m = _value_to_cpp_real_eigen(v, f"m{graph_id}") f = f"add_constant_probability_matrix(m{graph_id})" elif t is bn.ConstantSimplexMatrixNode: - m = _value_to_cpp_eigen(v, f"m{graph_id}") + m = _value_to_cpp_real_eigen(v, f"m{graph_id}") f = f"add_constant_col_simplex_matrix(m{graph_id})" elif t is bn.ConstantNaturalMatrixNode: - # TODO: This is wrong, it needs to be a MatrixXn - m = _value_to_cpp_eigen(v, f"m{graph_id}") + m = _value_to_cpp_natural_eigen(v, f"m{graph_id}") f = f"add_constant_natural_matrix(m{graph_id})" elif t is bn.ConstantBooleanMatrixNode: - # TODO: This is wrong, it needs to be a MatrixXb - m = _value_to_cpp_eigen(v, f"m{graph_id}") + m = _value_to_cpp_bool_eigen(v, f"m{graph_id}") f = f"add_constant_bool_matrix(m{graph_id})" else: f = "UNKNOWN" diff --git a/src/beanmachine/ppl/compiler/tests/bmg_query_test.py b/src/beanmachine/ppl/compiler/tests/bmg_query_test.py index 4a65939888..46969b7d50 100644 --- a/src/beanmachine/ppl/compiler/tests/bmg_query_test.py +++ b/src/beanmachine/ppl/compiler/tests/bmg_query_test.py @@ -118,9 +118,13 @@ def test_constant_functional(self) -> None: # Check if this is wrong and fix it. expected = """ graph::Graph g; -uint n0 = g.add_constant(torch::from_blob((float[]){2.5}, {})); +Eigen::MatrixXd m0(1, 1); +m0 << 2.5; +uint n0 = g.add_constant_real_matrix(m0); uint q0 = g.query(n0); -uint n1 = g.add_constant(torch::from_blob((float[]){1.5,-2.5}, {2})); +Eigen::MatrixXd m1(2, 1); +m1 << 1.5, -2.5; +uint n1 = g.add_constant_real_matrix(m1); uint q1 = g.query(n1); """ self.assertEqual(expected.strip(), observed.strip())