Added simd_extract_msbs intrinsic.

odin-lang · Nov 17, 2024 · 699d3ca · 699d3ca
1 parent 8f845c7
commit 699d3ca
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 0 deletions.
diff --git a/base/intrinsics/intrinsics.odin b/base/intrinsics/intrinsics.odin
@@ -283,6 +283,8 @@ simd_reduce_xor         :: proc(a: #simd[N]T) -> T where type_is_integer(T) || t
 simd_reduce_any         :: proc(a: #simd[N]T) -> T where type_is_boolean(T) ---
 simd_reduce_all         :: proc(a: #simd[N]T) -> T where type_is_boolean(T) ---
 
+simd_extract_msbs       :: proc(a: #simd[N]T) -> bit_set[0..<N] where type_is_integer(T) || type_is_boolean(T) ---
+
 
 simd_gather       :: proc(ptr: #simd[N]rawptr, val: #simd[N]T, mask: #simd[N]U) -> #simd[N]T where type_is_integer(U) || type_is_boolean(U) ---
 simd_scatter      :: proc(ptr: #simd[N]rawptr, val: #simd[N]T, mask: #simd[N]U)              where type_is_integer(U) || type_is_boolean(U) ---

diff --git a/core/simd/simd.odin b/core/simd/simd.odin
@@ -135,6 +135,8 @@ reduce_xor         :: intrinsics.simd_reduce_xor
 reduce_any         :: intrinsics.simd_reduce_any
 reduce_all         :: intrinsics.simd_reduce_all
 
+extract_msbs       :: intrinsics.simd_extract_msbs
+
 // swizzle :: proc(a: #simd[N]T, indices: ..int) -> #simd[len(indices)]T
 swizzle :: builtin.swizzle
 

diff --git a/src/check_builtin.cpp b/src/check_builtin.cpp
@@ -888,6 +888,38 @@ gb_internal bool check_builtin_simd_operation(CheckerContext *c, Operand *operan
 			return true;
 		}
 
+	case BuiltinProc_simd_extract_msbs:
+		{
+			Operand x = {};
+			check_expr(c, &x, ce->args[0]); if (x.mode == Addressing_Invalid) return false;
+
+			if (!is_type_simd_vector(x.type)) {
+				gbString xs = type_to_string(x.type);
+				error(x.expr, "'%.*s' expected a simd vector type, got '%s'", LIT(builtin_name), xs);
+				gb_string_free(xs);
+				return false;
+			}
+
+			Type *elem = base_array_type(x.type);
+			if (!is_type_integer_like(elem)) {
+				gbString xs = type_to_string(x.type);
+				error(x.expr, "'%.*s' expected a #simd type with integer or boolean elements, got '%s'", LIT(builtin_name), xs);
+				gb_string_free(xs);
+				return false;
+			}
+
+			i64 num_elems = get_array_type_count(x.type);
+
+			Type *result_type = alloc_type_bit_set();
+			result_type->BitSet.elem = t_int;
+			result_type->BitSet.lower = 0;
+			result_type->BitSet.upper = num_elems - 1;
+
+			operand->mode = Addressing_Value;
+			operand->type = result_type;
+			return true;
+		}
+
 
 	case BuiltinProc_simd_shuffle:
 		{

diff --git a/src/checker_builtin_procs.hpp b/src/checker_builtin_procs.hpp
@@ -181,6 +181,8 @@ BuiltinProc__simd_begin,
 	BuiltinProc_simd_reduce_any,
 	BuiltinProc_simd_reduce_all,
 
+	BuiltinProc_simd_extract_msbs,
+
 	BuiltinProc_simd_shuffle,
 	BuiltinProc_simd_select,
 
@@ -523,6 +525,8 @@ gb_global BuiltinProc builtin_procs[BuiltinProc_COUNT] = {
 	{STR_LIT("simd_reduce_any"),          1, false, Expr_Expr, BuiltinProcPkg_intrinsics},
 	{STR_LIT("simd_reduce_all"),         1, false, Expr_Expr, BuiltinProcPkg_intrinsics},
 
+	{STR_LIT("simd_extract_msbs"),       1, false, Expr_Expr, BuiltinProcPkg_intrinsics},
+
 
 	{STR_LIT("simd_shuffle"), 2, true,  Expr_Expr, BuiltinProcPkg_intrinsics},
 	{STR_LIT("simd_select"),  3, false, Expr_Expr, BuiltinProcPkg_intrinsics},

diff --git a/src/llvm_backend_proc.cpp b/src/llvm_backend_proc.cpp
@@ -1545,6 +1545,30 @@ gb_internal lbValue lb_build_builtin_simd_proc(lbProcedure *p, Ast *expr, TypeAn
 			return res;
 		}
 
+	case BuiltinProc_simd_extract_msbs:
+		{
+			Type *vt = arg0.type;
+			GB_ASSERT(vt->kind == Type_SimdVector);
+
+			i64 elem_bits = 8*type_size_of(elem);
+			i64 num_elems = get_array_type_count(vt);
+
+			LLVMTypeRef word_type = lb_type(m, elem);
+			LLVMValueRef shift_value = llvm_splat_int(num_elems, word_type, elem_bits - 1);
+			LLVMValueRef broadcast_value = LLVMBuildAShr(p->builder, arg0.value, shift_value, "");
+
+			LLVMTypeRef bitvec_type = LLVMVectorType(LLVMInt1TypeInContext(m->ctx), (unsigned)num_elems);
+			LLVMValueRef bitvec_value = LLVMBuildTrunc(p->builder, broadcast_value, bitvec_type, "");
+
+			LLVMTypeRef mask_type = LLVMIntTypeInContext(m->ctx, (unsigned)num_elems);
+			LLVMValueRef mask_value = LLVMBuildBitCast(p->builder, bitvec_value, mask_type, "");
+
+			LLVMTypeRef result_type = lb_type(m, res.type);
+			res.value = LLVMBuildZExtOrBitCast(p->builder, mask_value, result_type, "");
+
+			return res;
+		}
+
 
 	case BuiltinProc_simd_shuffle:
 		{