fix: make simd generic

Author: sstadick

ishlib/matcher/__init__.mojo

@@ -1,11 +1,29 @@
 from collections import Optional
 from gpu.host import DeviceBuffer
 from memory import Span, UnsafePointer
+from sys.info import sizeof, simdwidthof
+from sys.param_env import env_get_string
 
 from ishlib.matcher.alignment.scoring_matrix import MatrixKind
 from ishlib.searcher_settings import SemiGlobalEndsFreeness
 
 
+fn simd_width_selector[dtype: DType]() -> Int:
+    """Select a SIMD width for a given type.
+
+    This will override the `simdwidthof` if the env var ISH_SIMD_TARGET=baseline
+    """
+    constrained[sizeof[Scalar[dtype]]() <= 16, "dytpe size too large."]()
+    alias target = env_get_string["ISH_SIMD_TARGET", "none"]().lower()
+    alias SSE_U8_SIZE = 16
+
+    @parameter
+    if target == "baseline":
+        return SSE_U8_SIZE // sizeof[Scalar[dtype]]()
+    else:
+        return simdwidthof[dtype]()
+
+
 @value
 @register_passable("trivial")
 struct MatchResult:

ishlib/matcher/striped_local_matcher.mojo

@@ -5,7 +5,7 @@ from sys.info import (
 )
 
 from ishlib.vendor.log import Logger
-from ishlib.matcher import Matcher, MatchResult
+from ishlib.matcher import Matcher, MatchResult, simd_width_selector
 from ishlib.matcher.alignment import create_reversed
 from ishlib.matcher.alignment.scoring_matrix import ScoringMatrix, MatrixKind
 from ishlib.matcher.alignment.local_aln.striped import (
@@ -17,8 +17,8 @@ from ishlib.matcher.alignment.local_aln.striped import (
 
 @value
 struct StripedLocalMatcher[mut: Bool, //, origin: Origin[mut]](Matcher):
-    alias SIMD_U8_WIDTH = simdwidthof[UInt8]()
-    alias SIMD_U16_WIDTH = simdwidthof[UInt16]()
+    alias SIMD_U8_WIDTH = simd_width_selector[DType.uint8]()
+    alias SIMD_U16_WIDTH = simd_width_selector[DType.uint16]()
 
     var pattern: List[UInt8]
     var rev_pattern: List[UInt8]

ishlib/matcher/striped_semi_global_matcher.mojo

@@ -3,7 +3,7 @@ from gpu.host import DeviceContext
 from sys.info import simdwidthof
 
 from ishlib.gpu.kernels.semi_global import gpu_align_coarse
-from ishlib.matcher import Matcher, MatchResult
+from ishlib.matcher import Matcher, MatchResult, simd_width_selector
 from ishlib.matcher.alignment import create_reversed
 from ishlib.matcher.alignment.scoring_matrix import ScoringMatrix, MatrixKind
 from ishlib.matcher.alignment.semi_global_aln.striped import (
@@ -18,12 +18,9 @@ from ishlib.vendor.log import Logger
 
 @value
 struct StripedSemiGlobalMatcher(GpuMatcher):
-    alias SIMD_U8_WIDTH = simdwidthof[
-        UInt8
-    ]()  # // 4  # TODO: needs tuning on wider machines
-    alias SIMD_U16_WIDTH = simdwidthof[
-        UInt16
-    ]()  # // 4  # TODO: needs tuning on wider machines
+    alias SIMD_U8_WIDTH = simd_width_selector[DType.uint8]()
+    alias SIMD_U16_WIDTH = simd_width_selector[DType.uint16]()
+
     var pattern: List[UInt8]
     var rev_pattern: List[UInt8]
     var max_score: Int
@@ -46,6 +43,7 @@ struct StripedSemiGlobalMatcher(GpuMatcher):
         gap_open: UInt = 3,
         gap_extend: UInt = 1,
     ):
+        Logger.info("SIMD_U8_WIDTH=", Self.SIMD_U8_WIDTH)
         Logger.info("Performing matching with StripedSemiGlobalMatcher.")
         self.gap_open = gap_open
         self.gap_extend = gap_extend

pixi.toml

@@ -12,6 +12,7 @@ version = "1.1.0"
 
 [tasks]
 build = "pixi run mojo build -D ISH_LOG_LEVEL=info --march native -o ish main.mojo"
+build_generic = "pixi run mojo build -D ISH_SIMD_TARGET=baseline -D ISH_LOG_LEVEL=info -o ish main.mojo"
 debug = "pixi run mojo build -D ISH_LOG_LEVEL=debug --debug-level line-tables --march native -o ish main.mojo"
 pkglib = "pixi run mojo package ./ishlib"
 testlib = "pixi run mojo test tests -I . --co && pixi run mojo test tests -I ."