<commit_before>open Ir
open Analysis
open Util

let is_const_zero = function
  | IntImm 0 
  | FloatImm 0.0
  | Broadcast (IntImm 0, _)
  | Broadcast (FloatImm 0.0, _) -> true
  | _ -> false
and is_const_one = function
  | IntImm 1
  | FloatImm 1.0
  | Broadcast (IntImm 1, _)
  | Broadcast (FloatImm 1.0, _) -> true
  | _ -> false
and is_const = function
  | IntImm _ 
  | FloatImm _
  | Broadcast (IntImm _, _)
  | Broadcast (FloatImm _, _) -> true
  | _ -> false

(* Is an expression sufficiently simple that it should just be substituted in when it occurs in a let *)
let rec is_trivial = function
  | Var (_, _) -> true
  | Broadcast (Var (_, _), _) -> true
  | x -> is_const x

let rec is_simple x =
  (* One child simple, other children const *)
  let check = function
    | x when is_const x -> 0
    | x when is_simple x -> 1
    | _ -> 9999
  in
  let complexity = fold_children_in_expr check (+) 0 x in
  complexity <= 1

let rec constant_fold_expr expr = 
  let recurse = constant_fold_expr in

  if verbosity > 3 then dbg "Constant folding: %s\n" (Ir_printer.string_of_expr expr);
  
  match expr with
    (* Ignoring most const-casts for now, because we can't represent immediates of arbitrary types *)
    | Cast (t, e) when val_type_of_expr e = t -> recurse e
    | Cast (t, e) -> 
        begin match (t, recurse e) with
          | (Int 32, IntImm x)    -> IntImm x
          | (Int 32, FloatImm x)  -> IntImm (int_of_float x)
          | (Float 32, IntImm x)  -> FloatImm (float_of_int x)
          | (Float 32, FloatImm x) -> FloatImm x
          | (t, e)                -> Cast(t, recurse e)
        end

    (* basic binary ops *)
    | Bop (op, a, b) ->
      begin match (op, recurse a, recurse b) with
        | (_, IntImm   x, IntImm   y) -> IntImm   (caml_iop_of_bop op x y)
        | (_, FloatImm x, FloatImm y) -> FloatImm (caml_fop_of_bop op x y)

        (* Identity operations. These are not strictly constant
           folding, but they tend to come up at the same time *)
        | (Add, x, y) when is_const_zero x -> y
        | (Add, x, y) when is_const_zero y -> x
        | (Sub, x, y) when is_const_zero y -> x
        | (Sub, x, y) when x = y -> make_zero (val_type_of_expr x)
        | (Mul, x, y) when is_const_one x -> y
        | (Mul, x, y) when is_const_one y -> x
        | (Mul, x, y) when is_const_zero x -> x
        | (Mul, x, y) when is_const_zero y -> y
        | (Div, x, y) when is_const_one y -> x
        | (Div, x, y) when x = y -> make_one (val_type_of_expr x)

        (* Commonly occuring reducible integer div/mod operations:
           when m % n = 0
           x*m % n -> 0
           (x*m + y) % n -> (y % n)
           x*m / n -> x*(m/n)
           (x*m + y) / n -> x*(m/n) + y/n

           e.g.: (4*x + 1) % 2 -> 1
        *)
        | (Mod, Bop (Mul, IntImm m, x), IntImm n) 
        | (Mod, Bop (Mul, x, IntImm m), IntImm n) when (m mod n = 0) -> 
            IntImm 0
        | (Mod, Bop (Add, y, Bop (Mul, x, IntImm m)), IntImm n) 
        | (Mod, Bop (Add, y, Bop (Mul, IntImm m, x)), IntImm n) 
        | (Mod, Bop (Add, Bop (Mul, x, IntImm m), y), IntImm n) 
        | (Mod, Bop (Add, Bop (Mul, IntImm m, x), y), IntImm n) when (m mod n = 0) -> 
            recurse (y %~ (IntImm n))
        | (Mod, x, IntImm n) -> begin
          match reduce_expr_modulo x n with
            | Some k -> IntImm k
            | None -> Bop (Mod, x, IntImm n)
        end
        | (Div, Bop (Mul, IntImm m, x), IntImm n) 
        | (Div, Bop (Mul, x, IntImm m), IntImm n) when (m mod n = 0) -> 
            recurse (x *~ (IntImm (m/n)))
        | (Div, Bop (Add, y, Bop (Mul, x, IntImm m)), IntImm n) 
        | (Div, Bop (Add, y, Bop (Mul, IntImm m, x)), IntImm n) 
        | (Div, Bop (Add, Bop (Mul, x, IntImm m), y), IntImm n) 
        | (Div, Bop (Add, Bop (Mul, IntImm m, x), y), IntImm n) when (m mod n = 0) -> 
            recurse ((x *~ (IntImm (m/n))) +~ (y /~ (IntImm n)))


        (* op (Ramp, Broadcast) should be folded into the ramp *)
        | (Add, Broadcast (e, _), Ramp (b, s, n)) 
        | (Add, Ramp (b, s, n), Broadcast (e, _)) -> Ramp (recurse (b +~ e), s, n)
        | (Sub, Ramp (b, s, n), Broadcast (e, _)) -> Ramp (recurse (b -~ e), s, n)
        | (Mul, Broadcast (e, _), Ramp (b, s, n)) 
        | (Mul, Ramp (b, s, n), Broadcast (e, _)) -> Ramp (recurse (b *~ e), recurse (s *~ e), n)
        (* | (Div, Ramp (b, s, n), Broadcast (e, _)) -> Ramp (recurse (b /~ e), recurse (s /~ e), n) *)

        (* op (Broadcast, Broadcast) should be folded into the broadcast *)
        | (Add, Broadcast (a, n), Broadcast(b, _)) -> recurse (Broadcast (a +~ b, n))
        | (Sub, Broadcast (a, n), Broadcast(b, _)) -> recurse (Broadcast (a -~ b, n))
        | (Mul, Broadcast (a, n), Broadcast(b, _)) -> recurse (Broadcast (a *~ b, n))
        | (Div, Broadcast (a, n), Broadcast(b, _)) -> recurse (Broadcast (a /~ b, n))
        | (Mod, Broadcast (a, n), Broadcast(b, _)) -> recurse (Broadcast (a %~ b, n))
        | (Max, Broadcast (a, n), Broadcast(b, _)) -> recurse (Broadcast (Bop (Max, a, b), n))
        | (Min, Broadcast (a, n), Broadcast(b, _)) -> recurse (Broadcast (Bop (Min, a, b), n))

        (* Converting subtraction to addition *)
        | (Sub, x, IntImm y) -> recurse (x +~ (IntImm (-y)))
        | (Sub, x, FloatImm y) -> recurse (x +~ (FloatImm (-.y)))

        (* Convert const + varying to varying + const (reduces the number of cases to check later) *)
        | (Add, x, y) when is_const x -> recurse (y +~ x)
        | (Mul, x, y) when is_const x -> recurse (y *~ x)

        (* Convert divide by float constants to multiplication *)
        | (Div, x, FloatImm y) -> Bop (Mul, x, FloatImm (1.0 /. y))
        | (Div, x, Broadcast (FloatImm y, n)) -> Bop (Mul, x, Broadcast (FloatImm (1.0 /. y), n))

        (* Ternary expressions that can be reassociated. Previous passes have cut down on the number we need to check. *)
        (* (X + y) + z -> X + (y + z) *)
        | (Add, Bop (Add, x, y), z) when is_const y && is_const z -> recurse (x +~ (y +~ z))
        (* (X * y) * z -> X * (y * z) *)
        | (Mul, Bop (Mul, x, y), z) when is_const y && is_const z -> recurse (x *~ (y *~ z))
        (* (x - Y) + z -> (x + z) - Y *)
        | (Add, Bop (Sub, x, y), z) when is_const x && is_const z -> recurse ((x +~ z) -~ y)            

        (* In ternary expressions with one constant, pull the constant outside *)
        | (Add, Bop (Add, x, y), z) when is_const y -> recurse ((x +~ z) +~ y)
        | (Sub, Bop (Add, x, y), z) when is_const y -> recurse ((x -~ z) +~ y)
        | (Add, z, Bop (Add, x, y)) when is_const y -> recurse ((x +~ z) +~ y)
        | (Sub, z, Bop (Add, x, y)) when is_const y -> recurse ((z -~ x) -~ y)

        (* Additions or subtractions that cancel an inner term *)
        | (Sub, Bop (Add, x, y), z) 
        | (Add, z, Bop (Sub, x, y))
        | (Add, Bop (Sub, x, y), z) when y = z -> x
        | (Sub, Bop (Add, x, y), z) when x = z -> y

        | (Sub, Bop (Add, x, Bop (Add, y, z)), w) when w = z -> 
            recurse (x +~ y)
        | (Sub, Bop (Add, x, Bop (Add, z, y)), w) when w = z -> 
            recurse (x +~ y) 

        (* Ternary expressions that should be distributed *)
        | (Mul, Bop (Add, x, y), z) when is_const y && is_const z -> recurse ((x *~ z) +~ (y *~ z))     

        (* These particular patterns are commonly generated by lowering, so we should catch and simplify them *)
        | (Max, Bop (Add, x, IntImm a), Bop (Add, y, IntImm b)) when x = y ->
            if a > b then x +~ IntImm a else x +~ IntImm b              
        | (Max, x, Bop (Add, y, IntImm a)) 
        | (Max, Bop (Add, x, IntImm a), y) when x = y ->
            if a > 0 then x +~ IntImm a else x
        | (Min, Bop (Add, x, IntImm a), Bop (Add, y, IntImm b)) when x = y ->
            if a < b then x +~ IntImm a else x +~ IntImm b
        | (Min, x, Bop (Add, y, IntImm a)) 
        | (Min, Bop (Add, x, IntImm a), y) when x = y ->
            if a < 0 then x +~ IntImm a else x
        | (Min, Bop (Add, x, y), Bop (Add, y2, x2))
        | (Max, Bop (Add, x, y), Bop (Add, y2, x2)) when x = x2 && y = y2 ->
          x +~ y

        | (Min, x, y)
        | (Max, x, y) when x = y -> x


        | (op, x, y) -> Bop (op, x, y)
      end
<commit_msg>Constant folding does some factorization<commit_after>open Ir
open Analysis
open Util

let is_const_zero = function
  | IntImm 0 
  | FloatImm 0.0
  | Broadcast (IntImm 0, _)
  | Broadcast (FloatImm 0.0, _) -> true
  | _ -> false
and is_const_one = function
  | IntImm 1
  | FloatImm 1.0
  | Broadcast (IntImm 1, _)
  | Broadcast (FloatImm 1.0, _) -> true
  | _ -> false
and is_const = function
  | IntImm _ 
  | FloatImm _
  | Broadcast (IntImm _, _)
  | Broadcast (FloatImm _, _) -> true
  | _ -> false

(* Is an expression sufficiently simple that it should just be substituted in when it occurs in a let *)
let rec is_trivial = function
  | Var (_, _) -> true
  | Broadcast (Var (_, _), _) -> true
  | x -> is_const x

let rec is_simple x =
  (* One child simple, other children const *)
  let check = function
    | x when is_const x -> 0
    | x when is_simple x -> 1
    | _ -> 9999
  in
  let complexity = fold_children_in_expr check (+) 0 x in
  complexity <= 1

let rec constant_fold_expr expr = 
  let recurse = constant_fold_expr in

  if verbosity > 3 then dbg "Constant folding: %s\n" (Ir_printer.string_of_expr expr);
  
  match expr with
    (* Ignoring most const-casts for now, because we can't represent immediates of arbitrary types *)
    | Cast (t, e) when val_type_of_expr e = t -> recurse e
    | Cast (t, e) -> 
        begin match (t, recurse e) with
          | (Int 32, IntImm x)    -> IntImm x
          | (Int 32, FloatImm x)  -> IntImm (int_of_float x)
          | (Float 32, IntImm x)  -> FloatImm (float_of_int x)
          | (Float 32, FloatImm x) -> FloatImm x
          | (t, e)                -> Cast(t, recurse e)
        end

    (* basic binary ops *)
    | Bop (op, a, b) ->
      begin match (op, recurse a, recurse b) with
        | (_, IntImm   x, IntImm   y) -> IntImm   (caml_iop_of_bop op x y)
        | (_, FloatImm x, FloatImm y) -> FloatImm (caml_fop_of_bop op x y)

        (* Identity operations. These are not strictly constant
           folding, but they tend to come up at the same time *)
        | (Add, x, y) when is_const_zero x -> y
        | (Add, x, y) when is_const_zero y -> x
        | (Sub, x, y) when is_const_zero y -> x
        | (Sub, x, y) when x = y -> make_zero (val_type_of_expr x)
        | (Mul, x, y) when is_const_one x -> y
        | (Mul, x, y) when is_const_one y -> x
        | (Mul, x, y) when is_const_zero x -> x
        | (Mul, x, y) when is_const_zero y -> y
        | (Div, x, y) when is_const_zero x -> make_zero (val_type_of_expr x)
        | (Div, x, y) when is_const_one y -> x
        | (Div, x, y) when x = y -> make_one (val_type_of_expr x)

        (* Commonly occuring reducible integer div/mod operations:
           when m % n = 0
           x*m % n -> 0
           (x*m + y) % n -> (y % n)
           x*m / n -> x*(m/n)
           (x*m + y) / n -> x*(m/n) + y/n

           e.g.: (4*x + 1) % 2 -> 1
        *)
        | (Mod, Bop (Mul, IntImm m, x), IntImm n) 
        | (Mod, Bop (Mul, x, IntImm m), IntImm n) when (m mod n = 0) -> 
            IntImm 0
        | (Mod, Bop (Add, y, Bop (Mul, x, IntImm m)), IntImm n) 
        | (Mod, Bop (Add, y, Bop (Mul, IntImm m, x)), IntImm n) 
        | (Mod, Bop (Add, Bop (Mul, x, IntImm m), y), IntImm n) 
        | (Mod, Bop (Add, Bop (Mul, IntImm m, x), y), IntImm n) when (m mod n = 0) -> 
            recurse (y %~ (IntImm n))
        | (Mod, x, IntImm n) -> begin
          match reduce_expr_modulo x n with
            | Some k -> IntImm k
            | None -> Bop (Mod, x, IntImm n)
        end
        | (Div, Bop (Mul, IntImm m, x), IntImm n) 
        | (Div, Bop (Mul, x, IntImm m), IntImm n) when (m mod n = 0) -> 
            recurse (x *~ (IntImm (m/n)))
        | (Div, Bop (Add, y, Bop (Mul, x, IntImm m)), IntImm n) 
        | (Div, Bop (Add, y, Bop (Mul, IntImm m, x)), IntImm n) 
        | (Div, Bop (Add, Bop (Mul, x, IntImm m), y), IntImm n) 
        | (Div, Bop (Add, Bop (Mul, IntImm m, x), y), IntImm n) when (m mod n = 0) -> 
            recurse ((x *~ (IntImm (m/n))) +~ (y /~ (IntImm n)))


        (* op (Ramp, Broadcast) should be folded into the ramp *)
        | (Add, Broadcast (e, _), Ramp (b, s, n)) 
        | (Add, Ramp (b, s, n), Broadcast (e, _)) -> Ramp (recurse (b +~ e), s, n)
        | (Sub, Ramp (b, s, n), Broadcast (e, _)) -> Ramp (recurse (b -~ e), s, n)
        | (Mul, Broadcast (e, _), Ramp (b, s, n)) 
        | (Mul, Ramp (b, s, n), Broadcast (e, _)) -> Ramp (recurse (b *~ e), recurse (s *~ e), n)
        (* | (Div, Ramp (b, s, n), Broadcast (e, _)) -> Ramp (recurse (b /~ e), recurse (s /~ e), n) *)

        (* op (Broadcast, Broadcast) should be folded into the broadcast *)
        | (Add, Broadcast (a, n), Broadcast(b, _)) -> recurse (Broadcast (a +~ b, n))
        | (Sub, Broadcast (a, n), Broadcast(b, _)) -> recurse (Broadcast (a -~ b, n))
        | (Mul, Broadcast (a, n), Broadcast(b, _)) -> recurse (Broadcast (a *~ b, n))
        | (Div, Broadcast (a, n), Broadcast(b, _)) -> recurse (Broadcast (a /~ b, n))
        | (Mod, Broadcast (a, n), Broadcast(b, _)) -> recurse (Broadcast (a %~ b, n))
        | (Max, Broadcast (a, n), Broadcast(b, _)) -> recurse (Broadcast (Bop (Max, a, b), n))
        | (Min, Broadcast (a, n), Broadcast(b, _)) -> recurse (Broadcast (Bop (Min, a, b), n))

        (* Converting subtraction to addition *)
        | (Sub, x, IntImm y) -> recurse (x +~ (IntImm (-y)))
        | (Sub, x, FloatImm y) -> recurse (x +~ (FloatImm (-.y)))

        (* Convert const + varying to varying + const (reduces the number of cases to check later) *)
        | (Add, x, y) when is_const x -> recurse (y +~ x)
        | (Mul, x, y) when is_const x -> recurse (y *~ x)

        (* Convert divide by float constants to multiplication *)
        | (Div, x, FloatImm y) -> Bop (Mul, x, FloatImm (1.0 /. y))
        | (Div, x, Broadcast (FloatImm y, n)) -> Bop (Mul, x, Broadcast (FloatImm (1.0 /. y), n))

        (* Ternary expressions that can be reassociated. Previous passes have cut down on the number we need to check. *)
        (* (X + y) + z -> X + (y + z) *)
        | (Add, Bop (Add, x, y), z) when is_const y && is_const z -> recurse (x +~ (y +~ z))
        (* (X * y) * z -> X * (y * z) *)
        | (Mul, Bop (Mul, x, y), z) when is_const y && is_const z -> recurse (x *~ (y *~ z))
        (* (x - Y) + z -> (x + z) - Y *)
        | (Add, Bop (Sub, x, y), z) when is_const x && is_const z -> recurse ((x +~ z) -~ y)            

        (* In ternary expressions with one constant, pull the constant outside *)
        | (Add, Bop (Add, x, y), z) when is_const y -> recurse ((x +~ z) +~ y)
        | (Sub, Bop (Add, x, y), z) when is_const y -> recurse ((x -~ z) +~ y)
        | (Add, z, Bop (Add, x, y)) when is_const y -> recurse ((x +~ z) +~ y)
        | (Sub, z, Bop (Add, x, y)) when is_const y -> recurse ((z -~ x) -~ y)

        (* Additions or subtractions that cancel an inner term *)
        | (Sub, Bop (Add, x, y), z) 
        | (Add, z, Bop (Sub, x, y))
        | (Add, Bop (Sub, x, y), z) when y = z -> x
        | (Sub, Bop (Add, x, y), z) when x = z -> y

        | (Sub, Bop (Add, x, Bop (Add, y, z)), w) when w = z -> 
            recurse (x +~ y)
        | (Sub, Bop (Add, x, Bop (Add, z, y)), w) when w = z -> 
            recurse (x +~ y) 

        (* Ternary expressions that should be distributed *)
        | (Mul, Bop (Add, x, y), z) when is_const y && is_const z -> recurse ((x *~ z) +~ (y *~ z))     

        (* Pull out common factors *)
        | (Add, Bop (Mul, y1, x1), Bop (Mul, y2, x2)) 
        | (Add, Bop (Mul, x1, y1), Bop (Mul, y2, x2)) 
        | (Add, Bop (Mul, y1, x1), Bop (Mul, x2, y2)) 
        | (Add, Bop (Mul, x1, y1), Bop (Mul, x2, y2)) when y1 = y2 ->
          recurse ((x1 +~ x2) *~ y1)
            (* 
        | (Sub, Bop (Mul, y1, x1), Bop (Mul, y2, x2)) 
        | (Sub, Bop (Mul, x1, y1), Bop (Mul, y2, x2)) 
        | (Sub, Bop (Mul, y1, x1), Bop (Mul, x2, y2)) 
            *)
        | (Sub, Bop (Mul, x1, y1), Bop (Mul, x2, y2)) when y1 = y2 ->
          recurse ((x1 -~ x2) *~ y1)

        (* These particular patterns are commonly generated by lowering, so we should catch and simplify them *)
        | (Max, Bop (Add, x, IntImm a), Bop (Add, y, IntImm b)) when x = y ->
            if a > b then x +~ IntImm a else x +~ IntImm b              
        | (Max, x, Bop (Add, y, IntImm a)) 
        | (Max, Bop (Add, x, IntImm a), y) when x = y ->
            if a > 0 then x +~ IntImm a else x
        | (Min, Bop (Add, x, IntImm a), Bop (Add, y, IntImm b)) when x = y ->
            if a < b then x +~ IntImm a else x +~ IntImm b
        | (Min, x, Bop (Add, y, IntImm a)) 
        | (Min, Bop (Add, x, IntImm a), y) when x = y ->
            if a < 0 then x +~ IntImm a else x
        | (Min, Bop (Add, x, y), Bop (Add, y2, x2))
        | (Max, Bop (Add, x, y), Bop (Add, y2, x2)) when x = x2 && y = y2 ->
          x +~ y

        | (Min, x, y)
        | (Max, x, y) when x = y -> x

        | (op, x, y) -> Bop (op, x, y)
      end

{'fraction_non_alphanumeric': 0.16990487172095706, 'fraction_numerical': 0.008417411357739983, 'mean_word_length': 1.8982456140350876, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 178, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4783319', 'n_tokens_mistral': 6977, 'n_tokens_neox': 6629, 'n_words': 2441}