diff --git a/emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/shared/mma_exec.h b/emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/shared/mma_exec.h index eb21e415bec0..5a8eb6e0f88d 100644 --- a/emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/shared/mma_exec.h +++ b/emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/shared/mma_exec.h @@ -2334,7 +2334,9 @@ void exec_wmma_f32_f8_spec(auto &cu, uint32_t dst, uint32_t s0, uint32_t s1, uin } } -constexpr bool wmma_f32_f32_native_width_supported(uint32_t n, uint32_t width) { +/// Return whether an f32-accumulating matrix output row divides into complete +/// native SIMD chunks. +constexpr bool mma_f32_native_width_supported(uint32_t n, uint32_t width) { return width > 1 && n % width == 0; } @@ -2356,7 +2358,7 @@ void exec_wmma_f32_f32_spec(auto &cu, uint32_t dst, uint32_t s0, uint32_t s1, ui return; } else { constexpr uint32_t W = static_cast(util::native::size()); - if (util::force_scalar() || !wmma_f32_f32_native_width_supported(N, W)) { + if (util::force_scalar() || !mma_f32_native_width_supported(N, W)) { exec_wmma_f32(cu, M, N, K, in_bits, dst, s0, s1, s2, amdgpu::extract_f32, amdgpu::extract_f32, const_acc, c_modifier); return; @@ -4285,7 +4287,7 @@ void exec_f32_mfma_f32_spec(auto &cu, uint32_t dst, uint32_t s0, uint32_t s1, ui return; } else { constexpr uint32_t W = static_cast(util::native::size()); - if (util::force_scalar() || cbsz != 0 || blgp != 0 || W <= 1 || N % W != 0 || + if (util::force_scalar() || cbsz != 0 || blgp != 0 || !mma_f32_native_width_supported(N, W) || cu.wf_size() != 64) { exec_f32(cu, M, N, K, BATCH, in_bits, dst, s0, s1, s2, amdgpu::extract_f32, amdgpu::extract_f32, const_acc, cbsz, abid, blgp); @@ -4359,40 +4361,34 @@ void exec_f32_mfma_f32_spec(auto &cu, uint32_t dst, uint32_t s0, uint32_t s1, ui } } -/// Fast path for v_mfma_f32_16x16x32_f16. This single MFMA shape is the only -/// MFMA variant fired by OPT-125M fp16 eager forward (488k invocations per -/// forward; ~4B internal MACs). Kept as a dedicated specialization (rather -/// than forwarding to generic exec_f32) because the compile-time M/N/K/B let -/// the compiler fully unroll the 16-row x 32-K inner matmul into straight-line -/// AVX-512 FMAs — a runtime-dimension loop is materially slower on this hot -/// path. Hoists A and B into dense f32 buffers via extract_f16, runs the -/// matmul as 16 zmm-wide f32 FMA rows (512 zmm FMAs per MFMA). Batched shapes -/// stage every result on the stack and publish only after all operand reads, -/// preserving destructive-overlap semantics without a heap-backed Result -/// vector. VGPR access is batched through observed register-access regions -/// (one view per operand, no per-element virtual read_vgpr/write_vgpr). Falls -/// back to the generic exec_f32 when: +/// Fast path for the f16-input f32 MFMA shapes. The original target, +/// v_mfma_f32_16x16x32_f16, accounts for 488k invocations and about 4B +/// internal MACs per OPT-125M fp16 eager forward. Compile-time M/N/K/B let the +/// compiler fully unroll the inner matmul into straight-line native SIMD FMAs; +/// a runtime-dimension loop is materially slower on this hot path. The path +/// snapshots packed f16 inputs through observed register-access regions, +/// bulk-converts them into dense f32 buffers, and publishes staged results only +/// after every operand read. It falls back to the generic exec_f32 when: /// - is unavailable -/// - host native_simd is not 16 lanes (i.e. no AVX-512) +/// - host native_simd has no usable width for the output columns /// - cbsz/blgp lane permutation is non-default /// - RJ_FORCE_SCALAR is set template void exec_f32_mfma_f16_spec(auto &cu, uint32_t dst, uint32_t s0, uint32_t s1, uint32_t s2, uint32_t const_acc, uint32_t cbsz, uint32_t abid, uint32_t blgp) { constexpr uint32_t B = BATCH, in_bits = 16; - static_assert(N % 16 == 0, "specialized f16 MFMA assumes N is a multiple of the zmm width"); if constexpr (!util::has_stdx_simd) { exec_f32(cu, M, N, K, B, in_bits, dst, s0, s1, s2, amdgpu::extract_f16, amdgpu::extract_f16, const_acc, cbsz, abid, blgp); return; } else { - if (util::force_scalar() || cbsz != 0 || blgp != 0 || util::native::size() != 16 || + constexpr uint32_t W = static_cast(util::native::size()); + if (util::force_scalar() || cbsz != 0 || blgp != 0 || !mma_f32_native_width_supported(N, W) || cu.wf_size() != 64) { exec_f32(cu, M, N, K, B, in_bits, dst, s0, s1, s2, amdgpu::extract_f16, amdgpu::extract_f16, const_acc, cbsz, abid, blgp); return; } - constexpr uint32_t W = 16; // guaranteed by the native::size()==16 guard above const uint32_t wf = cu.wf_size(); auto reads = read_mfma_fast_path_regions(cu, s0, s1, s2, M, N, K, B, in_bits, const_acc, wf); auto writes = write_mfma_acc32_region(cu, dst, M, N, B, wf); @@ -4432,7 +4428,7 @@ void exec_f32_mfma_f16_spec(auto &cu, uint32_t dst, uint32_t s0, uint32_t s1, ui auto bl = input_loc(N, K, B, col, k, b, in_bits); B_buf[k * N + col] = B_f32[(bl.vgpr_offset * wf + bl.lane) * 2 + bl.sub_element]; } - // Dense MxKxN matmul, W-lane (zmm) stdx FMA over N (N/W chunks per row). + // Dense MxKxN matmul, W-lane native SIMD FMA over N (N/W chunks per row). for (uint32_t row = 0; row < M; ++row) for (uint32_t c0 = 0; c0 < N; c0 += W) { util::native c_row; @@ -4469,25 +4465,24 @@ void exec_f32_mfma_f16_spec(auto &cu, uint32_t dst, uint32_t s0, uint32_t s1, ui /// Fast path for the bf16-input MFMA shapes (v_mfma_f32_*_bf16). Identical to /// the f16 specialization except the bulk input convert is the bf16 zero-extend -/// (no F16C needed). Falls back to the generic exec_f32 without AVX-512 / under -/// force-scalar / with cbsz|blgp. +/// (no F16C needed). Falls back to the generic exec_f32 without usable host +/// SIMD, under force-scalar, or with cbsz/blgp. template void exec_f32_mfma_bf16_spec(auto &cu, uint32_t dst, uint32_t s0, uint32_t s1, uint32_t s2, uint32_t const_acc, uint32_t cbsz, uint32_t abid, uint32_t blgp) { constexpr uint32_t B = BATCH, in_bits = 16; - static_assert(N % 16 == 0, "specialized bf16 MFMA assumes N is a multiple of the zmm width"); if constexpr (!util::has_stdx_simd) { exec_f32(cu, M, N, K, B, in_bits, dst, s0, s1, s2, amdgpu::extract_bf16, amdgpu::extract_bf16, const_acc, cbsz, abid, blgp); return; } else { - if (util::force_scalar() || cbsz != 0 || blgp != 0 || util::native::size() != 16 || + constexpr uint32_t W = static_cast(util::native::size()); + if (util::force_scalar() || cbsz != 0 || blgp != 0 || !mma_f32_native_width_supported(N, W) || cu.wf_size() != 64) { exec_f32(cu, M, N, K, B, in_bits, dst, s0, s1, s2, amdgpu::extract_bf16, amdgpu::extract_bf16, const_acc, cbsz, abid, blgp); return; } - constexpr uint32_t W = 16; // guaranteed by the native::size()==16 guard above const uint32_t wf = cu.wf_size(); auto reads = read_mfma_fast_path_regions(cu, s0, s1, s2, M, N, K, B, in_bits, const_acc, wf); auto writes = write_mfma_acc32_region(cu, dst, M, N, B, wf); @@ -4525,7 +4520,7 @@ void exec_f32_mfma_bf16_spec(auto &cu, uint32_t dst, uint32_t s0, uint32_t s1, u auto bl = input_loc(N, K, B, col, k, b, in_bits); B_buf[k * N + col] = B_f32[(bl.vgpr_offset * wf + bl.lane) * 2 + bl.sub_element]; } - // Dense MxKxN matmul, W-lane (zmm) stdx FMA over N (N/W chunks per row). + // Dense MxKxN matmul, W-lane native SIMD FMA over N (N/W chunks per row). for (uint32_t row = 0; row < M; ++row) for (uint32_t c0 = 0; c0 < N; c0 += W) { util::native c_row; diff --git a/emulation/rocjitsu/tests/execution_plugin_test.cpp b/emulation/rocjitsu/tests/execution_plugin_test.cpp index a05148d01c67..aa01a9f16067 100644 --- a/emulation/rocjitsu/tests/execution_plugin_test.cpp +++ b/emulation/rocjitsu/tests/execution_plugin_test.cpp @@ -3347,11 +3347,11 @@ TEST(ExecutionPluginTest, WmmaReadObservationSkipsConstantAccumulator) { 0xFFFF'FFFFu); } -static_assert(!amdgpu::wmma_f32_f32_native_width_supported(16, 1)); -static_assert(amdgpu::wmma_f32_f32_native_width_supported(16, 4)); -static_assert(amdgpu::wmma_f32_f32_native_width_supported(16, 8)); -static_assert(amdgpu::wmma_f32_f32_native_width_supported(16, 16)); -static_assert(!amdgpu::wmma_f32_f32_native_width_supported(16, 32)); +static_assert(!amdgpu::mma_f32_native_width_supported(16, 1)); +static_assert(amdgpu::mma_f32_native_width_supported(16, 4)); +static_assert(amdgpu::mma_f32_native_width_supported(16, 8)); +static_assert(amdgpu::mma_f32_native_width_supported(16, 16)); +static_assert(!amdgpu::mma_f32_native_width_supported(16, 32)); TEST(ExecutionPluginTest, WmmaF32NativeWidthFastPathUsesRegionReads) { if constexpr (!util::has_stdx_simd) { @@ -3359,7 +3359,7 @@ TEST(ExecutionPluginTest, WmmaF32NativeWidthFastPathUsesRegionReads) { } else { constexpr uint32_t M = 16, N = 16, K = 4; constexpr uint32_t width = static_cast(util::native::size()); - if (!amdgpu::wmma_f32_f32_native_width_supported(N, width)) + if (!amdgpu::mma_f32_native_width_supported(N, width)) GTEST_SKIP() << "f32 WMMA shape is not divisible by the native SIMD width"; ForceScalarOverride force_simd(false); @@ -3424,7 +3424,7 @@ TEST(ExecutionPluginTest, MfmaF32NativeWidthFastPathUsesRegionReads) { } else { constexpr uint32_t M = 16, N = 16, K = 4, B = 1; constexpr uint32_t width = static_cast(util::native::size()); - if (width <= 1 || N % width != 0) + if (!amdgpu::mma_f32_native_width_supported(N, width)) GTEST_SKIP() << "f32 MFMA shape is not divisible by the native SIMD width"; ForceScalarOverride force_simd(false); @@ -3454,8 +3454,9 @@ TEST(ExecutionPluginTest, MfmaFastPathReadHookReportsRace) { if constexpr (!util::has_stdx_simd) { GTEST_SKIP() << "stdx SIMD is unavailable"; } else { - if (util::native::size() != 16) - GTEST_SKIP() << "MFMA fast path requires 16-lane native"; + constexpr uint32_t width = static_cast(util::native::size()); + if (!amdgpu::mma_f32_native_width_supported(16, width)) + GTEST_SKIP() << "f16 MFMA shape is not divisible by the native SIMD width"; struct ForceScalarGuard { bool old = util::force_scalar(); diff --git a/emulation/rocjitsu/tests/simd_correctness/mfma_lazy_storage_test.cpp b/emulation/rocjitsu/tests/simd_correctness/mfma_lazy_storage_test.cpp index dd79d05a0918..4f8b936eb47f 100644 --- a/emulation/rocjitsu/tests/simd_correctness/mfma_lazy_storage_test.cpp +++ b/emulation/rocjitsu/tests/simd_correctness/mfma_lazy_storage_test.cpp @@ -19,8 +19,9 @@ using MfmaLazyStorage = simdojo::detail::SoftwareLazyRegisterStorage::size() != 16) - GTEST_SKIP() << "test requires the 16-lane matrix fast path"; + constexpr uint32_t width = static_cast(util::native::size()); + if (!amdgpu::mma_f32_native_width_supported(32, width)) + GTEST_SKIP() << "f16 MFMA shape is not divisible by the native SIMD width"; constexpr uint32_t regs_per_chunk = MfmaLazyStorage::registers_per_chunk(); constexpr uint32_t source_a = regs_per_chunk - 1; diff --git a/emulation/rocjitsu/tests/simd_correctness/mfma_simd_exact_test.cpp b/emulation/rocjitsu/tests/simd_correctness/mfma_simd_exact_test.cpp index 3ffc81782696..694bcb4a42c4 100644 --- a/emulation/rocjitsu/tests/simd_correctness/mfma_simd_exact_test.cpp +++ b/emulation/rocjitsu/tests/simd_correctness/mfma_simd_exact_test.cpp @@ -171,6 +171,10 @@ TEST(MfmaSimdExact, Bf16Spec_AllShapes) { spec(amdgpu::exec_f32_mfma_bf16_spec<32, 32, 16>, "spec_bf16_32x32x16"); spec(amdgpu::exec_f32_mfma_bf16_spec<32, 32, 8>, "spec_bf16_32x32x8"); spec(amdgpu::exec_f32_mfma_bf16_spec<16, 16, 16>, "spec_bf16_16x16x16"); + spec(amdgpu::exec_f32_mfma_bf16_spec<32, 32, 4>, "spec_bf16_32x32x4"); + spec(amdgpu::exec_f32_mfma_bf16_spec<16, 16, 8>, "spec_bf16_16x16x8"); + spec(amdgpu::exec_f32_mfma_bf16_spec<32, 32, 2, 2>, "spec_bf16_32x32x2x2"); + spec(amdgpu::exec_f32_mfma_bf16_spec<16, 16, 2, 4>, "spec_bf16_16x16x2x4"); spec(amdgpu::exec_f32_mfma_bf16_spec<32, 32, 4, 2>, "spec_bf16_32x32x4x2"); spec(amdgpu::exec_f32_mfma_bf16_spec<16, 16, 4, 4>, "spec_bf16_16x16x4x4"); } @@ -239,8 +243,9 @@ TEST(MfmaSimdExact, F32SpecDestructiveSourceOverlap) { TEST(MfmaSimdExact, F16SpecDestructiveAccumulatorOverlap) { SKIP_IF_NO_SIMD(); - if (util::native::size() != 16) - GTEST_SKIP() << "test requires the 16-lane matrix fast path"; + constexpr uint32_t width = static_cast(util::native::size()); + if (!amdgpu::mma_f32_native_width_supported(32, width)) + GTEST_SKIP() << "f16 MFMA shape is not divisible by the native SIMD width"; constexpr uint32_t source_a = 0; constexpr uint32_t source_b = 16; constexpr uint32_t accumulator = 48; @@ -266,8 +271,9 @@ TEST(MfmaSimdExact, F16SpecDestructiveAccumulatorOverlap) { TEST(MfmaSimdExact, Bf16SpecDestructiveAccumulatorOverlap) { SKIP_IF_NO_SIMD(); - if (util::native::size() != 16) - GTEST_SKIP() << "test requires the 16-lane matrix fast path"; + constexpr uint32_t width = static_cast(util::native::size()); + if (!amdgpu::mma_f32_native_width_supported(32, width)) + GTEST_SKIP() << "bf16 MFMA shape is not divisible by the native SIMD width"; constexpr uint32_t source_a = 0; constexpr uint32_t source_b = 16; constexpr uint32_t accumulator = 48;