Postpone setting dim variables to their lower bounds till the very end.

lukstafi · lukstafi · commit 51bc961c154b · 2025-07-14T14:11:36.000+02:00
diff --git a/lib/row.ml b/lib/row.ml
@@ -664,7 +664,7 @@ let rec row_conjunction ?(id = phantom_row_id) stage constr1 constr2 =
                   elems_mismatch numerator (Num_elems known_product)
               else None))
 
-let rec apply_dim_constraint ~(source : source) ~(stage : stage) (dim : dim)
+let%track5_sexp rec apply_dim_constraint ~(source : source) ~(stage : stage) (dim : dim)
     (constr : dim_constraint) (env : environment) : constraint_ list * dim_constraint =
   let extras, constr =
     match (dim, constr) with
@@ -696,10 +696,12 @@ let rec apply_dim_constraint ~(source : source) ~(stage : stage) (dim : dim)
             | _ -> Option.value ~default:([], constr) @@ dim_conjunction constr bounds.constr))
     | _, Unconstrained_dim -> ([], constr)
   in
-  match (dim, constr, stage) with
+  (* FIXME: *)
+  (* match (dim, constr, stage) with
   | Var _, At_least_dim d, Stage4 ->
       (Dim_eq { d1 = dim; d2 = get_dim ~d () } :: extras, Unconstrained_dim)
-  | _ -> (extras, constr)
+  | _ -> *)
+  (extras, constr)
 
 exception Given_up
 
@@ -2012,7 +2014,7 @@ let%track5_sexp close_row_terminal ~(stage : stage) (env : environment)
           [%log6 "terminal row: keeping", (_r : row), "as", (r1 : row)];
           Terminal_row r1 :: term_dims ())
 
-let%debug5_sexp eliminate_dim_entry v ~lub constr =
+let%debug5_sexp eliminate_dim_entry ~final v ~lub constr =
   match (lub, constr) with
   | _, Unconstrained_dim | _, At_least_dim 1 -> None
   | Some (Dim { d; _ } as lub), At_least_dim d2 when d2 > d ->
@@ -2021,7 +2023,7 @@ let%debug5_sexp eliminate_dim_entry v ~lub constr =
            ( [%string "dereferenced at dimension %{d2#Int}, higher than use site"],
              [ Dim_mismatch [ lub; Var v ] ] )
   | Some lub, At_least_dim _ -> Some (Dim_eq { d1 = Var v; d2 = lub })
-  | None, At_least_dim d -> Some (Dim_eq { d1 = Var v; d2 = get_dim ~d () })
+  | None, At_least_dim d -> if final then Some (Dim_eq { d1 = Var v; d2 = get_dim ~d () }) else None
 
 let%debug5_sexp eliminate_variables (env : environment) ({ dims; bcast; id } as _r : row) :
     constraint_ list =
@@ -2030,8 +2032,8 @@ let%debug5_sexp eliminate_variables (env : environment) ({ dims; bcast; id } as
         Some
           (match Map.find env.dim_env v with
           | Some (Bounds_dim { lub; constr; _ }) ->
-              Option.value_or_thunk (eliminate_dim_entry v ~lub constr) ~default:(fun () ->
-                  Dim_eq { d1; d2 = get_dim ~d:1 () })
+              Option.value_or_thunk (eliminate_dim_entry ~final:true v ~lub constr)
+                ~default:(fun () -> Dim_eq { d1; d2 = get_dim ~d:1 () })
           | Some (Solved_dim _) -> assert false
           | None -> Dim_eq { d1; d2 = get_dim ~d:1 () })
     | _ -> None
@@ -2109,13 +2111,14 @@ let%debug4_sexp solve_inequalities ~(stage : stage) (ineqs : constraint_ list) (
   match stage with
   | Stage1 | Stage2 | Stage3 | Stage6 | Stage7 -> solve ineqs env
   | Stage4 ->
-      let finalize_lower_bound (v : dim_var) = function
-        | Bounds_dim { lub; constr; _ } -> Option.to_list @@ eliminate_dim_entry v ~lub constr
+      let finalize_upper_lower_bound (v : dim_var) = function
+        | Bounds_dim { lub; constr; _ } ->
+            Option.to_list @@ eliminate_dim_entry ~final:false v ~lub constr
         | _ -> []
       in
       let finalizing_entries : constraint_ list =
         Map.fold env.dim_env ~init:[] ~f:(fun ~key ~data accu ->
-            finalize_lower_bound key data @ accu)
+            finalize_upper_lower_bound key data @ accu)
       in
       solve (finalizing_entries @ ineqs) env
   | Stage5 ->
diff --git a/lib/shape_inference.md b/lib/shape_inference.md
@@ -163,10 +163,10 @@ The constraints are solved by: unification of the equation constraints, unificat
 * Stage 1 is online as tensors are composed, and conservatively performs unification and constraint propagation. Stages 2, 3, 4 are only performed once necessary: when projections or dimensions are requested.
 * Stage 2, when solving the constraints, substitutes dim variables in terminal shapes that do not have a LUB or other constraints, by dimension-1. (This is generalized at stage 6 to all variables.) (FIXME: reconsider this, see the algo for row variables: a new LUB can still be inferred.) Forces coefficients coming from precision byte sizes.
 * Stage 3, when solving the constraints, sets yet-unknown dimension and row variables in terminal shapes to their least upper bounds (if any). It substitutes row variables in terminal shapes that do not have a LUB by one axis if that's required to satisfy the variable's constraint.
-* Stage 4 sets yet-unknown dimensions with >1 lower bounds from direct accesses, to their LUBs if they have any, otherwise to the lower bound. It substitutes row variables in terminal shapes that do not have a LUB by no-further-axes. (This is generalized at stage 6 to all variables.)
+* Stage 4 sets yet-unknown dimensions with >1 lower bounds from direct accesses, to their LUBs if they have any. It substitutes row variables in terminal shapes that do not have a LUB by no-further-axes. (This is generalized at stage 6 to all variables.)
 * Stage 5 addresses `Total_elems` and `Exact` constraints with yet-unknown row variables. For `Total_elems` and a single row variable: if the constraint can be satisfied by assuming the row variable is no-further-axes, it sets the row variable to `Broadcastable`, otherwise it sets it to one axis of the required dimension. For multiple row variables, if one is of the Output kind, sets the other variables to no-further-axes, and retries.
 * Stage 6 sets row variables in the remaining inequalities to no-further-axes values. This can unlock further between-axis inequalities because of row variables sandwiched between leftmost axes from their side of the inequality and rightmost axes from the other side of the inequality.
-* Stage 7 sets all dim resp. row variables remaining in updated shapes to dimension-1 resp. no-further-axes.
+* Stage 7 sets all dim variables remaining in updated shapes to the lower bound if they have any, otherwise to dimension-1. It sets all row variables remaining in updated shapes to no-further-axes.
 
 Let's explain the shape inference functions.
 
