perf(webgpu): subgroup cooperative shaders for MatMul, BatchMatMul, Softmax (#141)

internal/backend/webgpu/backend.go

@@ -165,6 +165,11 @@
 		tensors map[*LazyGPUData]struct{}
 	}
 
+	// subgroupsEnabled is true when the device was created with the SubgroupOperations
+	// feature and subgroup shaders can be used safely. Detected once at device creation
+	// and used to select the subgroup MatMul path at dispatch time.
+	subgroupsEnabled bool
+
 	// Memory tracking
 	memoryStats struct {
 		totalAllocatedBytes uint64
@@ -229,7 +234,22 @@
 
 	info := adapter.Info()
 
-	device, err := adapter.RequestDevice(nil)
+	// Request SubgroupOperations when the adapter advertises support.
+	// On Vulkan the HAL ignores the feature flag at device creation (subgroup
+	// ops are governed by SPIR-V capabilities), but we still record whether
+	// the feature was granted so the rest of the backend can guard dispatch.
+	// On DX12 (SM 6.0+) and the Rust wgpu-native backend the flag is honored.
+	var desc *wgpu.DeviceDescriptor
+	subgroupsEnabled := false
+	if adapter.Features().Contains(gputypes.FeatureSubgroupOperations) {
+		desc = &wgpu.DeviceDescriptor{
+			Label:            "born",
+			RequiredFeatures: gputypes.Features(gputypes.FeatureSubgroupOperations),
+		}
+		subgroupsEnabled = true
+	}
+
+	device, err := adapter.RequestDevice(desc)
 	if err != nil {
 		adapter.Release()
 		instance.Release()
@@ -244,7 +264,12 @@
 		return nil, fmt.Errorf("webgpu: failed to get queue")
 	}
 
-	return newBackendFromDevice(instance, adapter, device, queue, &info)
+	b, err := newBackendFromDevice(instance, adapter, device, queue, &info)
+	if err != nil {
+		return nil, err
+	}
+	b.subgroupsEnabled = subgroupsEnabled
+	return b, nil
 }
 
 // newSoftwareBackend creates a Backend using the software HAL — CPU-based

internal/backend/webgpu/compute.go

@@ -656,9 +656,6 @@
 		return nil, fmt.Errorf("webgpu: matmul shape mismatch: [%d,%d] @ [%d,%d]", M, K, other.Shape()[0], N)
 	}
 
-	shader := b.compileShader("matmul", matmulShader)
-	entry := b.getOrCreatePipeline("matmul", shader, bglBinary)
-
 	bufferA := b.createBuffer(a.Data(), gputypes.BufferUsageStorage|gputypes.BufferUsageCopySrc)
 	defer bufferA.Release()
 
@@ -676,16 +673,35 @@
 	}
 	defer bufferResult.Release()
 
-	// Uniform: M, K, N as u32 (3×4 = 12 bytes, padded to 16).
-	params := make([]byte, 16)
-	binary.LittleEndian.PutUint32(params[0:4], M)
-	binary.LittleEndian.PutUint32(params[4:8], K)
-	binary.LittleEndian.PutUint32(params[8:12], N)
-	bufferParams := b.createUniformBuffer(params)
+	// Uniform buffer: M, K, N plus one u32 pad to meet std140 16-byte alignment.
+	paramBytes := make([]byte, 16)
+	binary.LittleEndian.PutUint32(paramBytes[0:4], M)
+	binary.LittleEndian.PutUint32(paramBytes[4:8], K)
+	binary.LittleEndian.PutUint32(paramBytes[8:12], N)
+	bufferParams := b.createUniformBuffer(paramBytes)
 	defer bufferParams.Release()
 
 	aSize := uint64(a.ByteSize())         //nolint:gosec // G115: integer overflow conversion int -> uint64
 	otherSize := uint64(other.ByteSize()) //nolint:gosec // G115: integer overflow conversion int -> uint64
+
+	var workgroupsX, workgroupsY uint32
+	var entry pipelineEntry
+
+	if b.subgroupsEnabled {
+		// Subgroup cooperative K-reduction: one workgroup(32 threads) per output element.
+		// Dispatch (N, M, 1) — col along X, row along Y.
+		shader := b.compileShader("matmul_subgroup", matmulSubgroupShader)
+		entry = b.getOrCreatePipeline("matmul_subgroup", shader, bglBinary)
+		workgroupsX = N
+		workgroupsY = M
+	} else {
+		// Scalar K-loop: 16×16 tile dispatch.
+		shader := b.compileShader("matmul", matmulShader)
+		entry = b.getOrCreatePipeline("matmul", shader, bglBinary)
+		workgroupsX = uint32(math.Ceil(float64(N) / 16.0))
+		workgroupsY = uint32(math.Ceil(float64(M) / 16.0))
+	}
+
 	bg := b.createBindGroupFromBuffers(entry.layout, []bindGroupBuffer{
 		bufBinding(bufferA, aSize),
 		bufBinding(bufferOther, otherSize),
@@ -694,9 +710,6 @@
 	})
 	defer bg.Release()
 
-	// 2D workgroup dispatch: 16×16 tiles.
-	workgroupsX := uint32(math.Ceil(float64(N) / 16.0))
-	workgroupsY := uint32(math.Ceil(float64(M) / 16.0))
 	resultData := b.execComputeAndRead(entry.pipeline, bg, workgroupsX, workgroupsY, 1, bufferResult, resultSize)
 
 	result, err := tensor.NewRaw(resultShape, tensor.Float32, tensor.WebGPU)
@@ -897,9 +910,6 @@
 
 	numClasses := uint32(input.Shape()[1]) //nolint:gosec // G115: safe, tensor dimensions are non-negative and fit in uint32
 
-	shader := b.compileShader("softmax", softmaxShader)
-	entry := b.getOrCreatePipeline("softmax", shader, bglUnary)
-
 	bufferInput := b.createBuffer(input.Data(), gputypes.BufferUsageStorage|gputypes.BufferUsageCopySrc)
 	defer bufferInput.Release()
 
@@ -913,22 +923,37 @@
 	}
 	defer bufferResult.Release()
 
-	// Uniform: batch_size, num_classes as u32.
-	params := make([]byte, 16)
-	binary.LittleEndian.PutUint32(params[0:4], batchSize)
-	binary.LittleEndian.PutUint32(params[4:8], numClasses)
-	bufferParams := b.createUniformBuffer(params)
+	// Uniform: batch_size, num_classes, two padding u32 to reach 16-byte std140 alignment.
+	paramBytes := make([]byte, 16)
+	binary.LittleEndian.PutUint32(paramBytes[0:4], batchSize)
+	binary.LittleEndian.PutUint32(paramBytes[4:8], numClasses)
+	bufferParams := b.createUniformBuffer(paramBytes)
 	defer bufferParams.Release()
 
+	var workgroups uint32
+	var entry pipelineEntry
+
+	if b.subgroupsEnabled {
+		// Subgroup cooperative reduction: one workgroup (32 lanes) per row.
+		// Dispatch (batchSize, 1, 1) — each workgroup handles exactly one row.
+		shader := b.compileShader("softmax_subgroup", softmaxSubgroupShader)
+		entry = b.getOrCreatePipeline("softmax_subgroup", shader, bglUnary)
+		workgroups = batchSize
+	} else {
+		// Scalar per-thread: one thread per row.
+		// Dispatch (ceil(batchSize/256), 1, 1).
+		shader := b.compileShader("softmax", softmaxShader)
+		entry = b.getOrCreatePipeline("softmax", shader, bglUnary)
+		workgroups = (batchSize + workgroupSize - 1) / workgroupSize
+	}
+
 	bg := b.createBindGroupFromBuffers(entry.layout, []bindGroupBuffer{
 		bufBinding(bufferInput, resultSize),
 		bufBinding(bufferResult, resultSize),
 		bufBinding(bufferParams, 16),
 	})
 	defer bg.Release()
 
-	// Each workgroup handles one row (batch sample).
-	workgroups := (batchSize + workgroupSize - 1) / workgroupSize
 	resultData := b.execComputeAndRead(entry.pipeline, bg, workgroups, 1, 1, bufferResult, resultSize)
 
 	result, err := tensor.NewRaw(input.Shape(), tensor.Float32, tensor.WebGPU)
@@ -984,9 +1009,6 @@
 		resultShape = tensor.Shape{shapeA[0], shapeA[1], int(M), int(N)}
 	}
 
-	shader := b.compileShader("batchMatMul", batchMatMulShader)
-	entry := b.getOrCreatePipeline("batchMatMul", shader, bglBinary)
-
 	bufferA := b.createBuffer(a.Data(), gputypes.BufferUsageStorage|gputypes.BufferUsageCopySrc)
 	defer bufferA.Release()
 
@@ -1003,16 +1025,35 @@
 	}
 	defer bufferResult.Release()
 
-	params := make([]byte, 16)
-	binary.LittleEndian.PutUint32(params[0:4], batch)
-	binary.LittleEndian.PutUint32(params[4:8], M)
-	binary.LittleEndian.PutUint32(params[8:12], K)
-	binary.LittleEndian.PutUint32(params[12:16], N)
-	bufferParams := b.createUniformBuffer(params)
+	paramBytes := make([]byte, 16)
+	binary.LittleEndian.PutUint32(paramBytes[0:4], batch)
+	binary.LittleEndian.PutUint32(paramBytes[4:8], M)
+	binary.LittleEndian.PutUint32(paramBytes[8:12], K)
+	binary.LittleEndian.PutUint32(paramBytes[12:16], N)
+	bufferParams := b.createUniformBuffer(paramBytes)
 	defer bufferParams.Release()
 
 	aSize := uint64(a.ByteSize())         //nolint:gosec // G115: integer overflow conversion int -> uint64
 	otherSize := uint64(other.ByteSize()) //nolint:gosec // G115: integer overflow conversion int -> uint64
+
+	var workgroupsX, workgroupsY uint32
+	var entry pipelineEntry
+
+	if b.subgroupsEnabled {
+		// Subgroup cooperative K-reduction: one workgroup(32 threads) per output element.
+		// Dispatch (N, M, batch) — col along X, row along Y, batch along Z.
+		shader := b.compileShader("batchMatMul_subgroup", batchMatMulSubgroupShader)
+		entry = b.getOrCreatePipeline("batchMatMul_subgroup", shader, bglBinary)
+		workgroupsX = N
+		workgroupsY = M
+	} else {
+		// Scalar K-loop: (N+7)/8 × (M+7)/8 × batch dispatch.
+		shader := b.compileShader("batchMatMul", batchMatMulShader)
+		entry = b.getOrCreatePipeline("batchMatMul", shader, bglBinary)
+		workgroupsX = (N + 7) / 8
+		workgroupsY = (M + 7) / 8
+	}
+
 	bg := b.createBindGroupFromBuffers(entry.layout, []bindGroupBuffer{
 		bufBinding(bufferA, aSize),
 		bufBinding(bufferB, otherSize),
@@ -1021,9 +1062,6 @@
 	})
 	defer bg.Release()
 
-	// Dispatch: (N+7)/8 × (M+7)/8 × batch.
-	workgroupsX := (N + 7) / 8
-	workgroupsY := (M + 7) / 8
 	resultData := b.execComputeAndRead(entry.pipeline, bg, workgroupsX, workgroupsY, batch, bufferResult, resultSize)
 
 	result, err := tensor.NewRaw(resultShape, tensor.Float32, tensor.WebGPU)