Exploit sharing in patricia tree combine operations #4073

bclement-ocp · 2025-06-02T09:47:45Z

This patch introduces a new union-like function to the patricia tree (and containers) interface, diff. It also provides _shared and _sharing variants of both diff and union to exploit internal sharing, and a typed interface to help ensure sharing is not accidentally lost.

The diff function is an union-like generalization of the existing diff_domains. diff f m1 m2 is similar to merge f m1 m2 with f None (Some _) = None and f (Some v) None = Some v.

The _shared variants (diff_shared, really) are the main motivation for this patch, and have a transformative impact on performance: diff_shared and union_shared exploit physical equality to simplify diff_shared f m m = empty and union_shared f m m = m in constant time. They also ensure maximal sharing of the result with their first argument.

The _sharing variants do not perform this optimisation but still ensure maximal sharing of the result with their first argument. They are useful to add or remove entries into a map while preserving sharing.

The Shared_set and Shared_map are thin wrapper modules that only expose sharing-preserving functions to avoid accidental loss of sharing.

These functions are currently unused, but will be used in the implementation #3219 and for the match-in-match heuristic (see #926), as discussed offline with @chambart and @lthls.

bclement-ocp · 2025-06-02T10:15:35Z

middle_end/flambda2/algorithms/container_types.mli

+      {b Complexity}: The complexity of [diff (union ss1 ss2) ss1] is linear
+      in the size of [ss2]. The complexity of [diff (add k ss) ss] is the
+      same as that of [mem k (add k ss)]. *)
+  val diff : t -> t -> t


I made this return a t because it felt more general, but in practice I end up always converting to a non-shared set / non-shared map.

I'll expand on the way these are intended to be used, because it might help clarify things. In the context where we are accumulating a "main" set/map (e.g. in a typing env), we use diff and diff_set / diff_map for two fundamentally different scenarios:

diff is used to efficiently compute the difference between an old version of the set/map and a new version of the set/map, in order to do some processing (e.g. a join) of the changes. It is crucial that sharing can be used to reduce the complexity (the size of the changes is typically small compared to the size of the main set), however sharing of the output is not important.

diff_set / diff_map is used to remove one (or more) elements from the main set/map; the result will be used as the new main set/map and sharing must be preserved.

Considering this, I think it makes more sense to return T.Set.t here, and probably also to make the _shared version not try to preserve the sharing of their output. If the reviewers agree, I'll make those changes.

This patch introduces a new `union`-like function to the patricia tree (and containers) interface, `diff`. It also provides `_shared` and `_sharing` variants of both `diff` and `union` to exploit internal sharing, and a typed interface to help ensure sharing is not accidentally lost. The `diff` function is an `union`-like generalization of the existing `diff_domains`. `diff f m1 m2` is similar to `merge f m1 m2` with `f None (Some _) = None` and `f (Some v) None = Some v`. The `_shared` variants (`diff_shared`, really) are the main motivation for this patch, and have a transformative impact on performance: `diff_shared` and `union_shared` exploit physical equality to simplify `diff_shared f m m = empty` and `union_shared f m m = m` in constant time. They also ensure maximal sharing of the result with their first argument. The `_sharing` variants do not perform this optimisation but still ensure maximal sharing of the result with their first argument. They are useful to add or remove entries into a map while preserving sharing. The `Shared_set` and `Shared_map` are thin wrapper modules that only expose sharing-preserving functions to avoid accidental loss of sharing.

@lthls

We currently deal with relational information through two (mostly, because the CSE environment adds backwards information to the typing environment) separate mechanisms: - The CSE environment keeps track of a map from arguments to values and effectively implements congruence closure over relations (i.e. it ensures that if `Get_tag(x) = y` and `Get_tag(x) = z` then `y = z`). - The typing env keeps track of a map from values to arguments though the relational types (if we have `Get_tag(x) = y` then `y` has type `(= Get_tag x)`) This is not ideal, for at least two reasons: the CSE environment and the canonicals in the typing env can get out of sync, causing missed optimisation opportunities such as oxcaml#3181; and the representation of relational types is restrictive (see e.g. oxcaml#2259, which this is can be considered a rewrite of). This patch introduces a new module in the flambda2 typing environment, the "types database". This module is intended to provide a single and principled source of relational informations between *canonical elements* (or equivalently of non-alias relational information). The design is loosely inspired from (a simplified version of) the description in oxcaml#3219. The types database internally maintains the same maps that the CSE environment and the typing environment maintained previously, except that it is wholly inside the typing environment now. In particular, the typing environment can make sure that the canonicals cannot get out of sync as with the CSE environment, and we also have access to the more precise n-way join when computing the join of the types database. Having a single source of truth for relational information, instead of it being distributed across the CSE environment and the typing environment, also makes it easier to implement a robust heuristic for match-in-match, which will come in a follow-up patch More precisely, the current patch implements the following: - The types database itself, with support for `Is_int`, `Get_tag` and `Is_null` relations ; - Integration in the typing env in order to properly maintain canonicals ; - Join of databases with the new n-way join algorithm Caveats: - This is currently gated behind the `-flambda2-types-database` command-line option or the `flambda2-types-database=1` OCAMLPARAM. This is done in order to facilitate testing, and we need to discuss during review whether to keep this flag or make it the default. - Enabling the types database currently enables the relational primitives but *does not* disable CSE for the corresponding primitives. This is because there are places (e.g. simplification of switch expressions) where we look up directly in the CSE environment, but these should be easy to replace with a lookup in the types database (it was not implemented in order to avoid multiplying the code paths, but will be implemented before merging). - Join is only implemented for the new n-way join algorithm. Implementing the join properly for the old binary join algorithm should not be complicated but requires a bit of thought and I'd like to wait until at least the first round of review first. - Support for env extension is currently partial. We discussed with @lthls on using identifiers for extensions in order to avoid potential infinite loops during meet, and this enables to store arbitrary information on extensions. Currently, the implementation is incorporating a database extension inside the `typing_env_extension` type, and allows attaching a full `typing_env_extension` to the database extension points. I believe that the first part is not necessary, and there are multiple places where database information is simply skipped in env extensions. Depending on the first round(s) of review, I'd be happy to either (a) keep things as is, or (b) more aggressively store database information in extensions (e.g. after a join) or (c) don't have database information in the `typing_env_extension` type at all (I personally favor (c); I believe this only has an effect on the join in practice and that we don't need to join such deep extensions). - Although the database supports `Untag_imm` and `Tag_imm` relations, they are not exposed through the typing env or used during simplification. This is in order to maintain feature parity with the existing implementation and facilitate comparisons; I plan on enabling them in a follow-up PR. Note: The implementation relies heavily on the shared patricia tree functions introduced in oxcaml#4073, and is rebased on oxcaml#4073; changes to the `middle_end/flambda2/algorithms/` directory in this PR do not need to be reviewed here.

@lthls

We currently deal with relational information through two (mostly, because the CSE environment adds backwards information to the typing environment) separate mechanisms: - The CSE environment keeps track of a map from arguments to values and effectively implements congruence closure over relations (i.e. it ensures that if `Get_tag(x) = y` and `Get_tag(x) = z` then `y = z`). - The typing env keeps track of a map from values to arguments though the relational types (if we have `Get_tag(x) = y` then `y` has type `(= Get_tag x)`) This is not ideal, for at least two reasons: the CSE environment and the canonicals in the typing env can get out of sync, causing missed optimisation opportunities such as oxcaml#3181; and the representation of relational types is restrictive (see e.g. oxcaml#2259, which this is can be considered a rewrite of). This patch introduces a new module in the flambda2 typing environment, the "types database". This module is intended to provide a single and principled source of relational informations between *canonical elements* (or equivalently of non-alias relational information). The design is loosely inspired from (a simplified version of) the description in oxcaml#3219. The types database internally maintains the same maps that the CSE environment and the typing environment maintained previously, except that it is wholly inside the typing environment now. In particular, the typing environment can make sure that the canonicals cannot get out of sync as with the CSE environment, and we also have access to the more precise n-way join when computing the join of the types database. Having a single source of truth for relational information, instead of it being distributed across the CSE environment and the typing environment, also makes it easier to implement a robust heuristic for match-in-match, which will come in a follow-up patch More precisely, the current patch implements the following: - The types database itself, with support for `Is_int`, `Get_tag` and `Is_null` relations ; - Integration in the typing env in order to properly maintain canonicals ; - Join of databases with the new n-way join algorithm Caveats: - This is currently gated behind the `-flambda2-types-database` command-line option or the `flambda2-types-database=1` OCAMLPARAM. This is done in order to facilitate testing, and we need to discuss during review whether to keep this flag or make it the default. - Enabling the types database currently enables the relational primitives but *does not* disable CSE for the corresponding primitives. This is because there are places (e.g. simplification of switch expressions) where we look up directly in the CSE environment, but these should be easy to replace with a lookup in the types database (it was not implemented in order to avoid multiplying the code paths, but will be implemented before merging). - Join is only implemented for the new n-way join algorithm. Implementing the join properly for the old binary join algorithm should not be complicated but requires a bit of thought and I'd like to wait until at least the first round of review first. - Support for env extension is currently partial. We discussed with @lthls on using identifiers for extensions in order to avoid potential infinite loops during meet, and this enables to store arbitrary information on extensions. Currently, the implementation is incorporating a database extension inside the `typing_env_extension` type, and allows attaching a full `typing_env_extension` to the database extension points. I believe that the first part is not necessary, and there are multiple places where database information is simply skipped in env extensions. Depending on the first round(s) of review, I'd be happy to either (a) keep things as is, or (b) more aggressively store database information in extensions (e.g. after a join) or (c) don't have database information in the `typing_env_extension` type at all (I personally favor (c); I believe this only has an effect on the join in practice and that we don't need to join such deep extensions). - Although the database supports `Untag_imm` and `Tag_imm` relations, they are not exposed through the typing env or used during simplification. This is in order to maintain feature parity with the existing implementation and facilitate comparisons; I plan on enabling them in a follow-up PR. Note: The implementation relies heavily on the shared patricia tree functions introduced in oxcaml#4073, and is rebased on oxcaml#4073; changes to the `middle_end/flambda2/algorithms/` directory in this PR do not need to be reviewed here.

bclement-ocp requested a review from chambart June 2, 2025 09:47

bclement-ocp added flambda2 Prerequisite for, or part of, flambda2 match-in-match prerequisites, or part of, match-in-match labels Jun 2, 2025

bclement-ocp commented Jun 2, 2025

View reviewed changes

bclement-ocp force-pushed the bclement/ptree-share branch from 6e9ede1 to 0d6ad09 Compare June 10, 2025 08:20

bclement-ocp mentioned this pull request Jun 10, 2025

Better treatment of relations in flambda2 types #4112

Open

chambart mentioned this pull request Jun 24, 2025

Sharing on patricia trees #4183

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Exploit sharing in patricia tree combine operations #4073

Exploit sharing in patricia tree combine operations #4073

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Exploit sharing in patricia tree combine operations #4073

Are you sure you want to change the base?

Exploit sharing in patricia tree combine operations #4073

Uh oh!

Conversation

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!