Skip to content

[WebAssembly] Add fold support for dot #151775

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
+109 −0
Open

[WebAssembly] Add fold support for dot #151775

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
4d304c8
Precommit test
badumbatish Aug 1, 2025
cb9aac0
Added combine support for dot
badumbatish Aug 1, 2025
6553946
Fix merge conflict from main
badumbatish Aug 1, 2025
86fe99b
Transition to tablegen for pattern
badumbatish Aug 5, 2025
34f58f1
Addresses PR reviews
badumbatish Aug 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions llvm/lib/Target/WebAssembly/WebAssemblyInstrSIMD.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1210,6 +1210,27 @@ defm EXTMUL_LOW_U :
defm EXTMUL_HIGH_U :
SIMDExtBinary<I64x2, extmul_high_u, "extmul_high_i32x4_u", 0xdf>;

// Pattern for dot
def : Pat<
(v4i32 (add
(wasm_shuffle
(v4i32 (extmul_low_s v8i16:$lhs, v8i16:$rhs)),
(v4i32 (extmul_high_s v8i16:$lhs, v8i16:$rhs)),
(i32 0), (i32 1), (i32 2), (i32 3),
(i32 8), (i32 9), (i32 10), (i32 11),
(i32 16), (i32 17), (i32 18), (i32 19),
(i32 24), (i32 25), (i32 26), (i32 27)),
(wasm_shuffle
(v4i32 (extmul_low_s v8i16:$lhs, v8i16:$rhs)),
(v4i32 (extmul_high_s v8i16:$lhs, v8i16:$rhs)),
(i32 4), (i32 5), (i32 6), (i32 7),
(i32 12), (i32 13), (i32 14), (i32 15),
(i32 20), (i32 21), (i32 22), (i32 23),
(i32 28), (i32 29), (i32 30), (i32 31)))
),
(v4i32 (DOT v8i16:$lhs, v8i16:$rhs))
>;

Comment on lines +1213 to +1233
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Move this pattern above so it's beside the other pattern for DOT?

//===----------------------------------------------------------------------===//
// Floating-point unary arithmetic
//===----------------------------------------------------------------------===//
Expand Down
88 changes: 88 additions & 0 deletions llvm/test/CodeGen/WebAssembly/simd-dot-reductions.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
; RUN: llc < %s -mattr=+simd128 | FileCheck %s

target triple = "wasm32-unknown-unknown"

define <4 x i32> @dot_sext_1(<8 x i16> %a, <8 x i16> %b) {
; CHECK-LABEL: dot_sext_1:
; CHECK: .functype dot_sext_1 (v128, v128) -> (v128)
; CHECK-NEXT: # %bb.0:
; CHECK-NEXT: local.get 0
; CHECK-NEXT: local.get 1
; CHECK-NEXT: i32x4.dot_i16x8_s
; CHECK-NEXT: # fallthrough-return
%sext1 = sext <8 x i16> %a to <8 x i32>
%sext2 = sext <8 x i16> %b to <8 x i32>
%mul = mul <8 x i32> %sext1, %sext2
%shuffle1 = shufflevector <8 x i32> %mul, <8 x i32> poison, <4 x i32> <i32 0, i32 2, i32 4, i32 6>
%shuffle2 = shufflevector <8 x i32> %mul, <8 x i32> poison, <4 x i32> <i32 1, i32 3, i32 5, i32 7>
%res = add <4 x i32> %shuffle1, %shuffle2
ret <4 x i32> %res
}


define <4 x i32> @dot_sext_2(<8 x i16> %a, <8 x i16> %b) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Out of curiosity can you also add a test case where we dot product one vector with itself? If it doesn't match that's ok, it doesn't need to be addressed in this PR

; CHECK-LABEL: dot_sext_2:
; CHECK: .functype dot_sext_2 (v128, v128) -> (v128)
; CHECK-NEXT: # %bb.0:
; CHECK-NEXT: local.get 0
; CHECK-NEXT: local.get 1
; CHECK-NEXT: i32x4.dot_i16x8_s
; CHECK-NEXT: # fallthrough-return
%sext1 = sext <8 x i16> %a to <8 x i32>
%sext2 = sext <8 x i16> %b to <8 x i32>
%mul = mul <8 x i32> %sext1, %sext2
%shuffle1 = shufflevector <8 x i32> %mul, <8 x i32> poison, <4 x i32> <i32 0, i32 2, i32 4, i32 6>
%shuffle2 = shufflevector <8 x i32> %mul, <8 x i32> poison, <4 x i32> <i32 1, i32 3, i32 5, i32 7>
%res = add <4 x i32> %shuffle2, %shuffle1
ret <4 x i32> %res
}

define <4 x i32> @dot_zext(<8 x i16> %a, <8 x i16> %b) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit, can you add a comment mentioning that this is a negative test

; CHECK-LABEL: dot_zext:
; CHECK: .functype dot_zext (v128, v128) -> (v128)
; CHECK-NEXT: .local v128
; CHECK-NEXT: # %bb.0:
; CHECK-NEXT: local.get 0
; CHECK-NEXT: local.get 1
; CHECK-NEXT: i32x4.extmul_low_i16x8_u
; CHECK-NEXT: local.tee 2
; CHECK-NEXT: local.get 0
; CHECK-NEXT: local.get 1
; CHECK-NEXT: i32x4.extmul_high_i16x8_u
; CHECK-NEXT: local.tee 1
; CHECK-NEXT: i8x16.shuffle 0, 1, 2, 3, 8, 9, 10, 11, 16, 17, 18, 19, 24, 25, 26, 27
; CHECK-NEXT: local.get 2
; CHECK-NEXT: local.get 1
; CHECK-NEXT: i8x16.shuffle 4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31
; CHECK-NEXT: i32x4.add
; CHECK-NEXT: # fallthrough-return
%zext1 = zext <8 x i16> %a to <8 x i32>
%zext2 = zext <8 x i16> %b to <8 x i32>
%mul = mul <8 x i32> %zext1, %zext2
%shuffle1 = shufflevector <8 x i32> %mul, <8 x i32> poison, <4 x i32> <i32 0, i32 2, i32 4, i32 6>
%shuffle2 = shufflevector <8 x i32> %mul, <8 x i32> poison, <4 x i32> <i32 1, i32 3, i32 5, i32 7>
%res = add <4 x i32> %shuffle1, %shuffle2
ret <4 x i32> %res
}

define <4 x i32> @dot_wrong_shuffle(<8 x i16> %a, <8 x i16> %b) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here, to comment that this is a negative test

; CHECK-LABEL: dot_wrong_shuffle:
; CHECK: .functype dot_wrong_shuffle (v128, v128) -> (v128)
; CHECK-NEXT: # %bb.0:
; CHECK-NEXT: local.get 0
; CHECK-NEXT: local.get 1
; CHECK-NEXT: i32x4.extmul_low_i16x8_s
; CHECK-NEXT: local.get 0
; CHECK-NEXT: local.get 1
; CHECK-NEXT: i32x4.extmul_high_i16x8_s
; CHECK-NEXT: i32x4.add
; CHECK-NEXT: # fallthrough-return
%sext1 = sext <8 x i16> %a to <8 x i32>
%sext2 = sext <8 x i16> %b to <8 x i32>
%mul = mul <8 x i32> %sext1, %sext2
%shuffle1 = shufflevector <8 x i32> %mul, <8 x i32> poison, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
%shuffle2 = shufflevector <8 x i32> %mul, <8 x i32> poison, <4 x i32> <i32 4, i32 5, i32 6, i32 7>
%res = add <4 x i32> %shuffle1, %shuffle2
ret <4 x i32> %res
}