Implement VectorElementSelection legalization for multi-precision vec…

…tors
nibrunie · Nov 18, 2018 · 8281af8 · 8281af8
1 parent 8aa7b99
commit 8281af8
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 2 deletions.
diff --git a/metalibm_core/core/ml_formats.py b/metalibm_core/core/ml_formats.py
@@ -1235,6 +1235,9 @@ def vector_format_builder(c_format_name, opencl_format_name, vector_size,
 v4trifloat64 = vector_format_builder("ml_tridouble4_t", "unsupported", 4, ML_TripleDouble, compound_constructor=ML_MultiPrecision_VectorFormat)
 v8trifloat64 = vector_format_builder("ml_tridouble8_t", "unsupported", 8, ML_TripleDouble, compound_constructor=ML_MultiPrecision_VectorFormat)
 
+LIST_SINGLE_MULTI_PRECISION_VECTOR_FORMATS = [v2dualfloat32, v3dualfloat32, v4dualfloat32, v8dualfloat32, v2trifloat32, v3trifloat32, v4trifloat32, v8trifloat32]
+LIST_DOUBLE_MULTI_PRECISION_VECTOR_FORMATS = [v2dualfloat64, v3dualfloat64, v4dualfloat64, v8dualfloat64, v2trifloat64, v3trifloat64, v4trifloat64, v8trifloat64]
+
 VECTOR_TYPE_MAP = {
     ML_Binary32: {
         2: v2float32,

diff --git a/metalibm_core/core/multi_precision.py b/metalibm_core/core/multi_precision.py
@@ -37,10 +37,13 @@
 ###############################################################################
 
 from metalibm_core.core.ml_operations import (
-    Comparison, LogicalAnd, LogicalOr,
+    Comparison, LogicalAnd, LogicalOr, BuildFromComponent,
+    VectorElementSelection,
 )
 from metalibm_core.core.ml_formats import ML_Bool
 
+from metalibm_core.opt.opt_utils import forward_attributes
+
 from metalibm_core.utility.log_report import Log
 
 
@@ -135,3 +138,31 @@ def legalize_mp_3elt_comparison(optree):
         )
     else:
         Log.report(Log.Error, "unsupported specifier {} in legalize_mp_2elt_comparison", specifier)
+
+
+def legalize_multi_precision_vector_element_selection(optree):
+    """ legalize a VectorElementSelection @p optree on a vector of
+        multi-precision elements to a single element """
+    assert isinstance(optree, VectorElementSelection)
+    multi_precision_vector = optree.get_input(0)
+    elt_index = optree.get_input(1)
+    hi_vector = multi_precision_vector.hi
+    lo_vector = multi_precision_vector.lo
+    limb_num = multi_precision_vector.get_precision().get_scalar_format().limb_num 
+    if limb_num == 2:
+        component_vectors = [hi_vector, lo_vector]
+    elif limb_num == 3:
+        me_vector = multi_precision_vector.me
+        component_vectors = [hi_vector, me_vector, lo_vector]
+    else:
+        Log.report(Log.Error, "unsupported number of limbs in legalize_multi_precision_vector_element_selection for {}", optree)
+    result = BuildFromComponent(
+        *tuple(VectorElementSelection(vector, elt_index) for vector in component_vectors)
+    )
+    forward_attributes(optree, result)
+    result.set_precision(optree.precision)
+    return result
+
+
+
+
diff --git a/metalibm_core/targets/common/vector_backend.py b/metalibm_core/targets/common/vector_backend.py
@@ -43,6 +43,9 @@
 from metalibm_core.core.legalizer import (
     min_legalizer, max_legalizer
 )
+from metalibm_core.core.multi_precision import (
+    legalize_multi_precision_vector_element_selection
+)
 
 from metalibm_core.code_generation.generator_utility import *
 from metalibm_core.code_generation.complex_generator import *
@@ -93,7 +96,11 @@ def legal_vector_element_selection(optree):
     optree.get_input(0).get_precision(),
     ML_CompoundVectorFormat
   )
-  return compound_format
+  multi_precision_scalar = isinstance(
+    optree.get_input(0).get_precision().get_scalar_format(),
+    ML_FP_MultiElementFormat
+  )
+  return compound_format and not multi_precision_scalar
 
 vector_opencl_code_generation_table = {
   BitLogicLeftShift: {
@@ -1256,6 +1263,12 @@ def unroll_vector_node(optree):
       legal_vector_element_selection: {
         lambda rformat, opformat, indexformat, optree: True: TemplateOperator("%s._[%s]", arity = 2, require_header = ["support_lib/ml_vector_format.h"]),
       },
+      lambda optree: not(legal_vector_element_selection(optree)): {
+          type_strict_match_list([ML_SingleSingle], LIST_SINGLE_MULTI_PRECISION_VECTOR_FORMATS, [ML_Integer]):
+            ComplexOperator(optree_modifier=legalize_multi_precision_vector_element_selection),
+          type_strict_match_list([ML_DoubleDouble], LIST_DOUBLE_MULTI_PRECISION_VECTOR_FORMATS, [ML_Integer]):
+            ComplexOperator(optree_modifier=legalize_multi_precision_vector_element_selection),
+      }
     },
   },
   LogicalNot: {