alpha 073-083

KunQuant/Driver.py

@@ -11,6 +11,8 @@ def optimize(f: Function, options: dict)->Dict[str, int]:
     if PassUtil.debug_mode:
         print("Before optimize: ", f)
     ret = infer_window(f, options)
+    # optimize before decompose to let value ranges work
+    special_optimize(f, options)
     decompose(f, options)
     expr_fold(f, options)
     special_optimize(f, options)

KunQuant/ops/CompOp.py

@@ -1,4 +1,4 @@
-from .ReduceOp import ReduceAdd, ReduceArgMax, ReduceRank, ReduceMin, ReduceMax, ReduceDecayLinear
+from .ReduceOp import ReduceAdd, ReduceMul, ReduceArgMax, ReduceRank, ReduceMin, ReduceMax, ReduceDecayLinear
 from KunQuant.Op import ConstantOp, OpBase, CompositiveOp, WindowedTrait, ForeachBackWindow, WindowedTempOutput, Builder, IterValue
 from .ElewiseOp import And, DivConst, GreaterThan, LessThan, Or, Select, SetInfOrNanToValue, Sub, Mul, Sqrt, SubConst, Div, CmpOp, Exp, Log
 from collections import OrderedDict
@@ -28,7 +28,11 @@ def decompose(self) -> List[OpBase]:
 class WindowedSum(WindowedReduce):
     def make_reduce(self, v: OpBase) -> OpBase:
         return ReduceAdd(v)
-
+
+class WindowedProduct(WindowedReduce):
+    def make_reduce(self, v: OpBase) -> OpBase:
+        return ReduceMul(v)
+
 class WindowedMin(WindowedReduce):
     def make_reduce(self, v: OpBase) -> OpBase:
         return ReduceMin(v)

KunQuant/ops/ElewiseOp.py

@@ -37,6 +37,9 @@ class And(BinaryElementwiseOp):
 class Max(BinaryElementwiseOp):
     pass
 
+class Min(BinaryElementwiseOp):
+    pass
+
 class Div(BinaryElementwiseOp):
     pass
 

KunQuant/ops/ReduceOp.py

@@ -3,6 +3,9 @@
 class ReduceAdd(ReductionOp):
     pass
 
+class ReduceMul(ReductionOp):
+    pass
+
 class ReduceMin(ReductionOp):
     pass
 

KunQuant/passes/Partitioner.py

@@ -2,6 +2,7 @@
 from KunQuant.Stage import Function, OpInfo
 from KunQuant.ops import GenericPartition
 from typing import List, Dict, Set, Tuple
+from .Util import debug_mode
 import typing
 from dataclasses import dataclass
 from collections import OrderedDict
@@ -58,7 +59,7 @@ def _collect_op_info(f: Function)-> Dict[OpBase, _PartitionOpInfo]:
         depender = set()
         for user in finfo.uses:
             depender.add(user)
-            depender.update(f.op_to_id[user].uses)
+            depender.update(ret[user].depender)
         ret[op] = _PartitionOpInfo(depender, dict([(inp, None) for inp in op.inputs]))
     return ret
 
@@ -197,6 +198,32 @@ def _search_output_use(op: OpBase, info: OpInfo) -> OpBase:
             return use
     return None
 
+def _print_partition_info(stages: List[GenericPartition], impl: List[Function]):
+    if debug_mode < 2:
+        return
+    name_to_impl = dict([(f.name, f) for f in impl])
+    for s in stages:
+        print(f'''===========================
+partition: {s.attrs["name"]}
+inputs: {[f.attrs["name"] for f in s.inputs]}
+Impl:
+{name_to_impl[s.attrs["name"]]}''')
+    if debug_mode < 3:
+        return
+    # loop check
+    for s in stages:
+        stck = []
+        def recurive(x: GenericPartition):
+            stck.append(x)
+            if x == s and len(stck) > 1:
+                raise RuntimeError("Loop in partitions: " + str([val.attrs["name"] for val in stck]))
+            for p in x.inputs:
+                recurive(p)
+            stck.pop()
+        recurive(s)
+
+
+
 def _transform_partitions(partitions: List[_Partition], f: Function) -> Tuple[Function, List[Function]]:
     naming_table = dict()
     for op in f.ops:
@@ -286,6 +313,7 @@ def add_to_naming_table(v: str) -> str:
         out_impl.append(p.impl_func)
         for dep in p.depending:
             p.stage_op.inputs.append(dep.stage_op)
+    _print_partition_info(out_stages, out_impl)
     out_stages = Function.topo_sort_ops(out_stages)
     return Function(out_stages), out_impl
 

KunQuant/passes/SpecialOpt.py

@@ -1,10 +1,44 @@
-from KunQuant.ops.MiscOp import BackRef
-from KunQuant.ops.ElewiseOp import AddConst, DivConst, Sqrt, Sub, Log, Div, CmpOp, Mul, Sign
+from KunQuant.ops import *
 from KunQuant.passes.Util import kun_pass
 from KunQuant.Op import Builder, OpBase, ForeachBackWindow, Rank, WindowedTempOutput, Output, IterValue, ConstantOp, Scale
-from KunQuant.ops import ReduceAdd, FastWindowedSum, SubConst, MulConst
 from KunQuant.Stage import Function
 from typing import List, Dict, Tuple
+from dataclasses import dataclass
+import math
+
+@dataclass
+class _ValueRange:
+    start: float
+    end: float
+    include_start: bool
+    include_end: bool
+    # def __lt__(self, other: '_ValueRange') -> bool:
+    #     return 
+
+# todo: more complex analysis
+class _ValueRangeManager:
+    def __init__(self) -> None:
+        self.values: Dict[OpBase, _ValueRange] = dict()
+
+    def _infer_range(self, op: OpBase) -> _ValueRange:
+        if op in self.values:
+            return self.values[op]
+        if isinstance(op, Rank):
+            return _ValueRange(0.0, 1.0, True, True)
+        if isinstance(op, TsRank):
+            return _ValueRange(0.0, op.attrs["window"]-1, True, True)
+        if isinstance(op, ConstantOp):
+            val = op.attrs["value"]
+            return _ValueRange(val, val, True, True)
+        return _ValueRange(-math.inf, math.inf, False, False)
+
+    def infer_range(self, op: OpBase) -> _ValueRange:
+        if op in self.values:
+            return self.values[op]
+        ret = self._infer_range(op)
+        self.values[op] = ret
+        return ret
+
 
 def _is_ok_for_reduce_opt(op: OpBase, enabled: bool) -> Tuple[OpBase, int]:
     if not enabled:
@@ -23,7 +57,7 @@ def _is_ok_for_reduce_opt(op: OpBase, enabled: bool) -> Tuple[OpBase, int]:
     window = loop.attrs["window"]
     return window_data, window
 
-def _is_sub_log(op: OpBase) -> OpBase:
+def _is_sub_log(ranges: _ValueRangeManager, op: OpBase) -> OpBase:
     '''
     if the op matches sub(log(X), log(Y)) or sub(log(X), backref(log(Y)))
     '''
@@ -48,7 +82,7 @@ def _is_sub_log(op: OpBase) -> OpBase:
 def _is_const_1(op: OpBase) -> bool:
     return isinstance(op, ConstantOp) and op.attrs["value"] == 1
 
-def _is_mul_1(op: OpBase) -> OpBase:
+def _is_mul_1(ranges: _ValueRangeManager, op: OpBase) -> OpBase:
     if isinstance(op, MulConst) and op.attrs["value"] == -1:
         return SubConst(op.inputs[0], 0, True)
     if isinstance(op, MulConst) and op.attrs["value"] == 1:
@@ -60,7 +94,7 @@ def _is_mul_1(op: OpBase) -> OpBase:
             return op.inputs[0]
     return None
 
-def _is_div_cmp_1(op: OpBase) -> OpBase:
+def _is_div_cmp_1(ranges: _ValueRangeManager, op: OpBase) -> OpBase:
     if isinstance(op, CmpOp):
         lhs, rhs = op.inputs
         if isinstance(lhs, Div):
@@ -80,7 +114,7 @@ def _is_div_cmp_1(op: OpBase) -> OpBase:
 _monotonic_ops = {Sqrt, Log, Scale, Rank}
 _monotonic_add = {AddConst, MulConst}
 _monotonic_sub = {SubConst, DivConst}
-def _is_rank_monotonic_inc(op: OpBase) -> OpBase:
+def _is_rank_monotonic_inc(ranges: _ValueRangeManager, op: OpBase) -> OpBase:
     '''
     check if is Rank(T(x)) where T is monotonicly increasing
     '''
@@ -93,9 +127,15 @@ def _is_rank_monotonic_inc(op: OpBase) -> OpBase:
         return Rank(internal.inputs[0])
     if internal.__class__ in _monotonic_add and internal.attrs["value"] > 0:
         return Rank(internal.inputs[0])
+    if isinstance(internal, Pow) and isinstance(internal.inputs[1], ConstantOp) and internal.inputs[1].attrs["value"] > 0:
+        # x^C, where C > 0. It is monotonic_inc when x > 0
+        base = internal.inputs[0]
+        ran = ranges.infer_range(base)
+        if ran.start >= 0:
+            return Rank(internal.inputs[0])
     return None
 
-def _is_sign_scale(op: OpBase) -> OpBase:
+def _is_sign_scale(ranges: _ValueRangeManager, op: OpBase) -> OpBase:
     '''
     check if is Sign(T(x)) where T does not change the signness of x
     '''
@@ -114,11 +154,12 @@ def special_impl(ops: List[OpBase], options: dict = {}) -> List[OpBase]:
     replace_map = dict()
     out = []
     changed = False
+    ranges = _ValueRangeManager()
     def _transform(f, op) -> bool:
         nonlocal changed, replace_map, out
         b = Builder(op.get_parent())
         with b:
-            newop = f(op)
+            newop = f(ranges, op)
         if newop is not None:
             if b.ops:
                 out.extend(b.ops)

KunQuant/predefined/Alpha101.py

@@ -55,6 +55,9 @@ def covariance(v: OpBase, v2: OpBase, window: int) -> OpBase:
 def sma(v: OpBase, window: int) -> OpBase:
     return WindowedAvg(v, window)
 
+def bool_to_10(v: OpBase) -> OpBase:
+    return Select(v, ConstantOp(1), ConstantOp(0))
+
 scale = Scale
 
 delay = BackRef
@@ -371,46 +374,46 @@ def alpha060(self: AllData):
 # Alpha#61	 (rank((vwap - ts_min(vwap, 16.1219))) < rank(correlation(vwap, adv180, 17.9282)))
 def alpha061(self: AllData):
     adv180 = sma(self.volume, 180)
-    return Select(rank((self.vwap - ts_min(self.vwap, 16))) < rank(correlation(self.vwap, adv180, 18)), ConstantOp(1), ConstantOp(0))
+    return bool_to_10(rank((self.vwap - ts_min(self.vwap, 16))) < rank(correlation(self.vwap, adv180, 18)))
 
 # Alpha#62	 ((rank(correlation(vwap, sum(adv20, 22.4101), 9.91009)) < rank(((rank(open) +rank(open)) < (rank(((high + low) / 2)) + rank(high))))) * -1)
 def alpha062(self: AllData):
     adv20 = sma(self.volume, 20)
     v1 = rank(correlation(self.vwap, sma(adv20, 22), 10))
     v2 = (rank(self.open) +rank(self.open))
     v3 = (rank(((self.high + self.low) / 2)) + rank(self.high))
-    v4 = Select(v2 < v3, ConstantOp(1), ConstantOp(0))
-    v5 = Select(v1 < rank(v4), ConstantOp(1), ConstantOp(0))
+    v4 = bool_to_10(v2 < v3)
+    v5 = bool_to_10(v1 < rank(v4))
     return (v5 * -1)
 
 # Alpha#64	 ((rank(correlation(sum(((open * 0.178404) + (low * (1 - 0.178404))), 12.7054),sum(adv120, 12.7054), 16.6208)) < rank(delta(((((high + low) / 2) * 0.178404) + (vwap * (1 -0.178404))), 3.69741))) * -1)
 def alpha064(self: AllData):
     adv120 = sma(self.volume, 120)
     a = rank(correlation(sma(((self.open * 0.178404) + (self.low * (1 - 0.178404))), 13),sma(adv120, 13), 17))
     b = rank(delta(((((self.high + self.low) / 2) * 0.178404) + (self.vwap * (1 -0.178404))), 4))
-    c = Select(a < b, ConstantOp(1), ConstantOp(0))
+    c = bool_to_10(a < b)
     return (c * -1)
 
 # Alpha#65	 ((rank(correlation(((open * 0.00817205) + (vwap * (1 - 0.00817205))), sum(adv60,8.6911), 6.40374)) < rank((open - ts_min(open, 13.635)))) * -1)
 def alpha065(self: AllData):
     adv60 = sma(self.volume, 60)
     a = rank(correlation(((self.open * 0.00817205) + (self.vwap * (1 - 0.00817205))), sma(adv60,9), 6))
     b = rank((self.open - ts_min(self.open, 14)))
-    return (Select(a < b, ConstantOp(1), ConstantOp(0)) * -1)
+    return (bool_to_10(a < b) * -1)
 
 # Alpha#66	 ((rank(decay_linear(delta(vwap, 3.51013), 7.23052)) + Ts_Rank(decay_linear(((((low* 0.96633) + (low * (1 - 0.96633))) - vwap) / (open - ((high + low) / 2))), 11.4157), 6.72611)) * -1)
 def alpha066(self: AllData):
     return ((rank(decay_linear(delta(self.vwap, 4), 7)) + ts_rank(decay_linear(((((self.low* 0.96633) + (self.low * (1 - 0.96633))) - self.vwap) / (self.open - ((self.high + self.low) / 2))), 11), 7)) * -1)
 
 # Alpha#68	 ((Ts_Rank(correlation(rank(high), rank(adv15), 8.91644), 13.9333) <rank(delta(((close * 0.518371) + (low * (1 - 0.518371))), 1.06157))) * -1)
-def alpha068(self):
+def alpha068(self: AllData):
     adv15 = sma(self.volume, 15)
     a = ts_rank(correlation(rank(self.high), rank(adv15), 9), 14)
     b = rank(delta(((self.close * 0.518371) + (self.low * (1 - 0.518371))), 1))
-    return (Select(a < b, ConstantOp(1), ConstantOp(0)) * -1)
+    return (bool_to_10(a < b) * -1)
 
 # Alpha#71	 max(Ts_Rank(decay_linear(correlation(Ts_Rank(close, 3.43976), Ts_Rank(adv180,12.0647), 18.0175), 4.20501), 15.6948), Ts_Rank(decay_linear((rank(((low + open) - (vwap +vwap)))^2), 16.4662), 4.4388))
-def alpha071(self):
+def alpha071(self: AllData):
     adv180 = sma(self.volume, 180)
     p1=ts_rank(decay_linear(correlation(ts_rank(self.close, 3), ts_rank(adv180,12), 18), 4), 16)
     inner = Pow(rank(((self.low + self.open) - (self.vwap +self.vwap))), ConstantOp(2))
@@ -419,16 +422,61 @@ def alpha071(self):
     #return max(ts_rank(decay_linear(correlation(ts_rank(self.close, 3), ts_rank(adv180,12), 18).to_frame(), 4).CLOSE, 16), ts_rank(decay_linear((rank(((self.low + self.open) - (self.vwap +self.vwap))).pow(2)).to_frame(), 16).CLOSE, 4))
 
 # Alpha#72	 (rank(decay_linear(correlation(((high + low) / 2), adv40, 8.93345), 10.1519)) /rank(decay_linear(correlation(Ts_Rank(vwap, 3.72469), Ts_Rank(volume, 18.5188), 6.86671),2.95011)))
-def alpha072(self):
+def alpha072(self: AllData):
     adv40 = sma(self.volume, 40)
     a = rank(decay_linear(correlation(((self.high + self.low) / 2), adv40, 9), 10)) + 0.0001
     b = rank(decay_linear(correlation(ts_rank(self.vwap, 4), ts_rank(self.volume, 19), 7),3)) + 0.0001
     return (a / b)    
 
+# Alpha#73	 (max(rank(decay_linear(delta(vwap, 4.72775), 2.91864)),Ts_Rank(decay_linear(((delta(((open * 0.147155) + (low * (1 - 0.147155))), 2.03608) / ((open *0.147155) + (low * (1 - 0.147155)))) * -1), 3.33829), 16.7411)) * -1)
+def alpha073(self: AllData):
+    p1=rank(decay_linear(delta(self.vwap, 5), 3))
+    p2=ts_rank(decay_linear(((delta(((self.open * 0.147155) + (self.low * (1 - 0.147155))), 2) / ((self.open *0.147155) + (self.low * (1 - 0.147155)))) * -1), 3), 17)
+    return -1* Max(p1, p2)
+
+# Alpha#74	 ((rank(correlation(close, sum(adv30, 37.4843), 15.1365)) <rank(correlation(rank(((high * 0.0261661) + (vwap * (1 - 0.0261661)))), rank(volume), 11.4791)))* -1)
+def alpha074(self: AllData):
+    adv30 = sma(self.volume, 30)
+    a = rank(correlation(self.close, sma(adv30, 37), 15))
+    b = rank(correlation(rank(((self.high * 0.0261661) + (self.vwap * (1 - 0.0261661)))), rank(self.volume), 11))
+    return (bool_to_10(a < b)* -1)
+
+# Alpha#75	 (rank(correlation(vwap, volume, 4.24304)) < rank(correlation(rank(low), rank(adv50),12.4413)))
+def alpha075(self: AllData):
+    adv50 = sma(self.volume, 50)
+    return bool_to_10(rank(correlation(self.vwap, self.volume, 4)) < rank(correlation(rank(self.low), rank(adv50),12)))
+
+def alpha077(self: AllData):
+    adv40 = sma(self.volume, 40)
+    p1=rank(decay_linear(((((self.high + self.low) / 2) + self.high) - (self.vwap + self.high)), 20))
+    p2=rank(decay_linear(correlation(((self.high + self.low) / 2), adv40, 3), 6))
+    return Min(p1, p2)
+
+# Alpha#78	 (rank(correlation(sum(((low * 0.352233) + (vwap * (1 - 0.352233))), 19.7428),sum(adv40, 19.7428), 6.83313))^rank(correlation(rank(vwap), rank(volume), 5.77492)))
+def alpha078(self: AllData):
+    adv40 = sma(self.volume, 40)
+    a = rank(correlation(ts_sum(((self.low * 0.352233) + (self.vwap * (1 - 0.352233))), 20),ts_sum(adv40,20), 7))
+    b = rank(SetInfOrNanToValue(correlation(rank(self.vwap), rank(self.volume), 6)))
+    return Pow(a, b)
+
+# Alpha#81	 ((rank(Log(product(rank((rank(correlation(vwap, sum(adv10, 49.6054),8.47743))^4)), 14.9655))) < rank(correlation(rank(vwap), rank(volume), 5.07914))) * -1)
+def alpha081(self: AllData):
+    adv10 = sma(self.volume, 10)
+    inner = Pow(rank(correlation(self.vwap, ts_sum(adv10, 50),8)), ConstantOp(4))
+    a = rank(Log(WindowedProduct(rank(inner), 15)))
+    b = rank(correlation(rank(self.vwap), rank(self.volume), 5))
+    return (bool_to_10(a < b) * -1)
+
+# Alpha#83	 ((rank(delay(((high - low) / (sum(close, 5) / 5)), 2)) * rank(rank(volume))) / (((high -low) / (sum(close, 5) / 5)) / (vwap - close)))
+def alpha083(self: AllData):
+    return ((rank(delay(((self.high - self.low) / (ts_sum(self.close, 5) / 5)), 2)) * rank(rank(self.volume))) / (((self.high -self.low) / (ts_sum(self.close, 5) / 5)) / (self.vwap - self.close)))
+
+
 all_alpha = [alpha001, alpha002, alpha003, alpha004, alpha005, alpha006, alpha007, alpha008, alpha009, alpha010,
     alpha011, alpha012, alpha013, alpha014, alpha015, alpha016, alpha017, alpha018, alpha019, alpha020, alpha021,
     alpha022, alpha023, alpha024, alpha025, alpha026, alpha027, alpha028, alpha029, alpha030, alpha031, alpha032,
     alpha033, alpha034, alpha035, alpha036, alpha037, alpha038, alpha039, alpha040, alpha041, alpha042, alpha043,
     alpha044, alpha045, alpha046, alpha047, alpha049, alpha050, alpha051, alpha052, alpha053, alpha054, alpha055,
-    alpha057, alpha060, alpha061, alpha062, alpha064, alpha065, alpha066, alpha068, alpha071, alpha072
+    alpha057, alpha060, alpha061, alpha062, alpha064, alpha065, alpha066, alpha068, alpha071, alpha072, alpha073,
+    alpha074, alpha075, alpha077, alpha078, alpha081, alpha083
     ]

cpp/Kun/Ops.hpp

@@ -168,6 +168,12 @@ struct ReduceAdd {
     operator f32x8() { return v; }
 };
 
+struct ReduceMul {
+    f32x8 v = _mm256_setzero_ps();
+    void step(f32x8 input, size_t index) { v = _mm256_mul_ps(v, input); }
+    operator f32x8() { return v; }
+};
+
 struct ReduceMin {
     f32x8 v = _mm256_set1_ps(std::numeric_limits<float>::infinity());
     void step(f32x8 input, size_t index) { v = _mm256_min_ps(v, input); }
@@ -262,7 +268,11 @@ struct ReduceRank {
 };
 
 inline f32x8 Max(f32x8 a, f32x8 b) {
-    return kun_simd::sc_max(a, b);
+    return _mm256_max_ps(a, b);
+}
+
+inline f32x8 Min(f32x8 a, f32x8 b) {
+    return _mm256_min_ps(a, b);
 }
 
 inline f32x8 Abs(f32x8 a) { return kun_simd::sc_abs(kun_simd::vec_f32x8(a)); }