optimize abs

KunQuant/passes/SpecialOpt.py

@@ -57,6 +57,17 @@ def _is_ok_for_reduce_opt(op: OpBase, enabled: bool) -> Tuple[OpBase, int]:
     window = loop.attrs["window"]
     return window_data, window
 
+def _is_abs_positive(ranges: _ValueRangeManager, op: OpBase) -> OpBase:
+    '''
+    if the op matches abs(X), where X>=0
+    '''
+    if not isinstance(op, Abs):
+        return None
+    inner = op.inputs[0]
+    if ranges.infer_range(inner).start >= 0:
+        return inner
+    return None
+
 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)))
@@ -183,6 +194,8 @@ def _transform(f, op) -> bool:
             continue
         if _transform(_is_sign_scale, op):
             continue
+        if _transform(_is_abs_positive, op):
+            continue
 
         # if it is reduce-sum in non-loop context
         result = _is_ok_for_reduce_opt(op, options.get("opt_reduce", True))

tests/test_alpha101.py

@@ -142,6 +142,7 @@ def make_data_and_ref(num_stock, num_time, ischeck):
     "rtol" : {
         "alpha013": 0.1,
         "alpha017": 0.1,
+        "alpha018": 0.05,
         "alpha014": 1e-2,
         "alpha016": 0.1,
         "alpha027": 0.1,
@@ -173,7 +174,7 @@ def make_data_and_ref(num_stock, num_time, ischeck):
         "alpha074": 0.04,
         "alpha075": 0.09,
         "alpha077": 0.001,
-        "alpha078": 0.001,
+        "alpha078": 0.01,
         "alpha081": 0.27,
     }
 }