type dim_cmp = Dim_var.comparator_witness

type dim_var_set = Set.M(Dim_var).t [@@deriving equal, sexp]

type 'a dim_map = 'a Map.M(Dim_var).t [@@deriving equal, sexp]

let dim_var_set_empty = Set.empty (module Dim_var)

let dim_map_empty = Map.empty (module Dim_var)

let solved_dim_to_string style { d; label; proj_id } =

match proj_id with

| None -> (

match style with

| `Only_labels -> ( match label with None -> "_" | Some l -> l)

else

match value with

| Dim s -> ((coeff_of_v, s) :: solved, new_unsolved)

| Affine { solved = value_solved; unsolved = value_unsolved } ->

(* Inlining affine expression: coeff_of_v * value *)

let scaled_solved = List.map value_solved ~f:(fun (c, sd) -> (coeff_of_v * c, sd)) in

let fresh_row_proj r =

let fresh_dim = function

| Var _ as d -> d

| Affine _ as d -> d (* FIXME: Affine dimensions don't get fresh projections *)

in

[@@deriving compare, equal, sexp]

type error_trace += Projection_mismatch of proj list

Sexp.List (Sexp.Atom "Projection_mismatch" :: List.map ps ~f:sexp_of_proj)

| error_trace -> sexp_of_error_trace error_trace

type proj_env = {

proj_to_index : proj_to_index;

proj_classes : proj_classes;

}

[@@deriving sexp]

| d -> (

match subst_dim env d with

| Dim { proj_id = Some proj_id; d; label = _ } -> Proj { proj_id; d }

| Var v when Map.mem proj_axis_env v -> Solved (Map.find_exn proj_axis_env v)

| Var v -> Var v

| Affine _ -> failwith "get_proj_equations: affine dimensions not supported in projections")

let v_env = dim_hashtbl () in

let p_solved = ref [] in

let p_dims = ref [] in

let rec loop = function

p_dims := (p1, d) :: !p_dims

| Proj_eq (Var v1, Var v2) when equal_dim_var v1 v2 -> ()

| Proj_eq ((Proj { proj_id = p1; d = d1 } as proj1), (Proj { proj_id = p2; d = d2 } as proj2))

( "Conflicting dimensions for the same projection",

[ Projection_mismatch [ proj1; proj2 ] ] );

p_dims := (p1, d1) :: !p_dims;

| Proj_eq (Proj p, Solved idx) | Proj_eq (Solved idx, Proj p) ->

p_solved := (p.proj_id, idx) :: !p_solved

| Proj_eq (Solved idx1, Solved idx2) when Idx.equal_axis_index idx1 idx2 -> ()

| Proj_eq (Solved idx1, Solved idx2) ->

raise

| Some p2 -> loop (Proj_eq (p, p2)))

| Iterated (Solved _) -> ()

| Iterated (Proj { proj_id; d }) -> p_dims := (proj_id, d) :: !p_dims

| Iterated (Var v) -> (

match Hashtbl.find v_env v with

| None ->

let idx = Idx.(Iterator (get_symbol ())) in

Hashtbl.add_exn v_env ~key:v ~data:(Solved idx)

| Some (Var v2) -> loop (Iterated (Var v2))

| Some (Solved _) -> ()

in

List.iter eqs ~f:loop;

List.iter !p_solved ~f:(fun (p, idx) ->

non_product := Set.add !non_product repr;

Utils.mref_add projs ~key:repr ~data:idx ~or_:(fun idx2 ->

if not @@ Idx.equal_axis_index idx idx2 then

raise

@@ Shape_error

("Multiple constraints on the same projection", [ Index_mismatch [ idx; idx2 ] ])));

List.iter !p_dims ~f:(fun (p, d) ->

if Idx.iterated d && (not @@ Map.mem !projs repr) then

Utils.mref_add product_dim ~key:repr ~data:d ~or_:(fun d2 ->

if d <> d2 then

Projection_mismatch [ Proj { proj_id = p; d }; Proj { proj_id = p; d = d2 } ];

] )));

Map.iteri !product_dim ~f:(fun ~key:p ~data:_ ->

Utils.mref_add_missing projs repr ~f:(fun () -> Idx.(Iterator (get_symbol ()))));

{

proj_classes = !proj_classes;

non_product = !non_product;

}

| Dim { d; _ } when not @@ Idx.iterated d -> Idx.Fixed_idx 0

| Dim { proj_id = None; _ } -> assert false

| Var v as dim ->

^ Sexp.to_string_hum ([%sexp_of: dim_var] v),

[ Dim_mismatch [ dim ] ] )

| Dim { proj_id = Some proj_id; d; _ } -> (

match Map.find proj_env.proj_to_index repr with

| Some i -> i

let proj_repr proj_env p =

let get_product_proj proj_env dim =

match dim with

val solve_proj_equations : proj_equation list -> proj_env

val get_proj_index : proj_env -> dim -> Ir.Indexing.axis_index

let rhs_dims = Array.of_list_map ~f:to_dims rhs in

let all_dims : Row.dim list = List.concat_map ~f:dims_of @@ (lhs :: rhs) in

(* Note: the ordering will affect performance of naive backends. *)

@@ List.filter_map all_dims ~f:(Row.get_product_proj proj_env)

in

let product_space : int array = Array.of_list_map all_product_projs ~f:snd in

* `get_proj_equations inequalities proj_axis_env env` converts both equations and inequalitites to projection equations. For inequalities, it takes broadcasting into account, and equates a potentially-broadcasted dim-1 projection to `Fixed_idx 0`. `proj_axis_env` originates from the `Shape` module, holds projections from the slice operator and the einsum syntax.

* `solve_proj_equations` unifies the projection equations, using union-find to maintain a representative for equal projections. Projections that already have an `axis_index` are `non_product` (not to be iterated over). The remaining projections have a `product_dim`, and get a fresh iterator.