diff --git a/tools/clang/unittests/HLSLExec/LinAlgTests.cpp b/tools/clang/unittests/HLSLExec/LinAlgTests.cpp index 1c23ec733c..43fce827c3 100644 --- a/tools/clang/unittests/HLSLExec/LinAlgTests.cpp +++ b/tools/clang/unittests/HLSLExec/LinAlgTests.cpp @@ -3152,6 +3152,9 @@ class DxilConf_SM610_LinAlg { TEST_METHOD(MatMatMul_Wave_16x16x16_I32); TEST_METHOD(MatMatMulAccum_Wave_16x16x16_F16); TEST_METHOD(MatMatMulAccum_Wave_8x32x16_F16_ToF32_NonUniform); + TEST_METHOD(MatMatMul_ThreadGroup_8x16x8_F16_NonUniform); + TEST_METHOD(MatMatMulAccum_ThreadGroup_8x16x8_F16_ToF32_NonUniform); + TEST_METHOD(MatMatMul_ThreadGroup_8x8x8_I32); TEST_METHOD(MatAccum_Wave_16x16_F16); TEST_METHOD(MatAccum_Wave_8x32_F16_BUse_NonUniform); @@ -5518,6 +5521,155 @@ static HRESULT selectWaveArithmeticMultiplyWaveSize( return S_OK; } +static UINT selectThreadGroupMatMulSize( + const linalg_test::ThreadGroupMatrixMultiplySupport &Support, + UINT WaveSize) { + constexpr UINT MaxThreadsPerGroup = + D3D12_CS_THREAD_GROUP_MAX_THREADS_PER_GROUP; + if (!Support.supported() || WaveSize == 0) + return 0; + + if (Support.PreferredThreadGroupSize > WaveSize && + Support.PreferredThreadGroupSize <= MaxThreadsPerGroup && + Support.supportsThreadGroupSize(Support.PreferredThreadGroupSize)) + return Support.PreferredThreadGroupSize; + + const UINT MaxThreadGroupSize = + std::min(Support.MaxThreadGroupSize, MaxThreadsPerGroup); + for (UINT ThreadGroupSize = Support.MinThreadGroupSize; + ThreadGroupSize <= MaxThreadGroupSize; + ThreadGroupSize += Support.MinThreadGroupSize) { + if (ThreadGroupSize > WaveSize) + return ThreadGroupSize; + } + + if (Support.PreferredThreadGroupSize <= MaxThreadGroupSize && + Support.supportsThreadGroupSize(Support.PreferredThreadGroupSize)) + return Support.PreferredThreadGroupSize; + if (Support.MinThreadGroupSize <= MaxThreadGroupSize) + return Support.MinThreadGroupSize; + return 0; +} + +static HRESULT selectThreadGroupMatMulConfiguration( + ID3D12Device *Device, const MatrixMultiplyCase &Case, LPCWSTR CaseName, + bool &Supported, UINT &SelectedWaveSize, UINT &SelectedThreadGroupSize) { + Supported = false; + SelectedWaveSize = 0; + SelectedThreadGroupSize = 0; + if (!Device) + return E_INVALIDARG; + if (!CaseName) + return E_INVALIDARG; + + const linalg_abi::D3D12_LINEAR_ALGEBRA_DATATYPE MatrixAType = + *toCapabilityDataType(Case.MatrixAType); + const linalg_abi::D3D12_LINEAR_ALGEBRA_DATATYPE MatrixBType = + *toCapabilityDataType(Case.MatrixBType); + const linalg_abi::D3D12_LINEAR_ALGEBRA_DATATYPE AccumulatorType = + *toCapabilityDataType(Case.AccumulatorType); + + // The HRESULT reports whether the capability queries themselves ran; the + // Supported out-parameter reports whether a usable configuration was found. + // Leaving Supported false and returning success is how a case is routed to a + // capability-gated skip instead of a failure. + linalg_test::TierSupport Tier; + HRESULT HR = linalg_test::queryTierSupport(Device, Tier); + VERIFY_SUCCEEDED(HR, "Linear algebra tier query must succeed"); + if (!Tier.supported()) + return S_OK; + + UINT MinWaveSize = 0; + UINT MaxWaveSize = 0; + HR = queryLaunchableWaveSizes(Device, MinWaveSize, MaxWaveSize); + VERIFY_SUCCEEDED(HR, "Launchable wave size query must succeed"); + if (MinWaveSize == 0) { + hlsl_test::LogCommentFmt( + L"Wave operations are unsupported; ThreadGroupMatrixMultiply is not " + L"applicable"); + return S_OK; + } + + const linalg_abi::D3D12_LINEAR_ALGEBRA_MATRIX_SHAPE Shape = { + Case.M, + Case.K, + Case.N, + }; + + for (UINT WaveSize = 4; WaveSize <= 128; WaveSize *= 2) { + if (WaveSize < MinWaveSize || WaveSize > MaxWaveSize) + continue; + + bool RolesConstructible = false; + HR = matrixMultiplyRolesConstructible(Device, Case, WaveSize, MatrixAType, + MatrixBType, AccumulatorType, + RolesConstructible); + VERIFY_SUCCEEDED(HR, "Matrix role construction query must succeed"); + if (!RolesConstructible) + continue; + + linalg_test::ThreadGroupMatrixMultiplySupport Multiply; + HR = linalg_test::queryThreadGroupMatrixMultiply( + Device, {{WaveSize, MatrixAType, MatrixBType, AccumulatorType}, Shape}, + Multiply); + VERIFY_SUCCEEDED(HR, "ThreadGroup matrix multiply query must succeed"); + if (!Multiply.supported()) + continue; + + const UINT ThreadGroupSize = + selectThreadGroupMatMulSize(Multiply, WaveSize); + if (ThreadGroupSize == 0) { + hlsl_test::LogCommentFmt( + L"ThreadGroupMatrixMultiply supports %s at wave=%u, but no legal " + L"shader group size is available: min=%u, max=%u, preferred=%u", + CaseName, WaveSize, Multiply.MinThreadGroupSize, + Multiply.MaxThreadGroupSize, Multiply.PreferredThreadGroupSize); + continue; + } + + hlsl_test::LogCommentFmt( + L"ThreadGroup matrix arithmetic capability matched wave=%u, " + L"threads=%u, crossWave=%u, shape=(%u,%u,%u) for %s", + WaveSize, ThreadGroupSize, ThreadGroupSize > WaveSize, Case.M, Case.K, + Case.N, CaseName); + Supported = true; + SelectedWaveSize = WaveSize; + SelectedThreadGroupSize = ThreadGroupSize; + return S_OK; + } + + hlsl_test::LogCommentFmt( + L"No executable ThreadGroupMatrixMultiply configuration supports %s", + CaseName); + // Every query succeeded and none reported support, so Supported stays false + // and the case is skipped rather than failed. + return S_OK; +} + +static bool threadGroupMatMulApplicable(ID3D12Device *Device, + const MatrixMultiplyCase &Case, + LPCWSTR CaseName, + UINT &SelectedWaveSize, + UINT &SelectedThreadGroupSize) { + bool Supported = false; + const HRESULT QueryResult = selectThreadGroupMatMulConfiguration( + Device, Case, CaseName, Supported, SelectedWaveSize, + SelectedThreadGroupSize); + if (!applyApplicability( + linalg_test::classifyApplicability( + QueryResult, Supported, + linalg_test::CapabilityRequirement::CapabilityGated), + CaseName)) + return false; + + VERIFY_IS_TRUE(SelectedWaveSize != 0, + "A case cleared to run must have a selected wave size"); + VERIFY_IS_TRUE( + SelectedThreadGroupSize != 0, + "A ThreadGroup case cleared to run must have a selected group size"); + return true; +} + static const char MatrixMultiplyShader[] = R"( #define USE_A 0 #define USE_B 1 @@ -5577,6 +5729,8 @@ static std::optional buildMatrixMultiplyCompilerArgs(const MatrixMultiplyCase &Case, MatrixScope Scope, UINT WaveSize, UINT NumThreads) { + if (Scope != MatrixScope::Wave && Scope != MatrixScope::ThreadGroup) + return std::nullopt; if (WaveSize == 0 || NumThreads == 0) return std::nullopt; @@ -5823,6 +5977,88 @@ void DxilConf_SM610_LinAlg::MatMatMul_Wave_16x16x16_I32() { L"MatMatMul_Wave_16x16x16_I32", VerboseLogging); } +static void runThreadGroupMultiplyCase(ID3D12Device *Device, + dxc::SpecificDllLoader &DxcSupport, + const MatrixMultiplyCase &Case, + LPCWSTR CaseName, bool Verbose) { + VERIFY_IS_TRUE(isMatrixMultiplyCaseValid(Case)); + if (!isMatrixMultiplyCaseValid(Case)) + return; + + UINT SelectedWaveSize = 0; + UINT SelectedThreadGroupSize = 0; + if (!threadGroupMatMulApplicable(Device, Case, CaseName, SelectedWaveSize, + SelectedThreadGroupSize)) + return; + + runMatrixMultiplyCase(Device, DxcSupport, Case, MatrixScope::ThreadGroup, + SelectedWaveSize, SelectedThreadGroupSize, Verbose); +} + +static MatrixMultiplyCase +makeRectangularF16ThreadGroupMultiplyCase(ComponentType AccumulatorType, + MatrixMultiplyOperation Operation) { + MatrixMultiplyCase Case = {}; + Case.MatrixAType = ComponentType::F16; + Case.MatrixBType = ComponentType::F16; + Case.AccumulatorType = AccumulatorType; + Case.M = 8; + Case.K = 16; + Case.N = 8; + Case.Operation = Operation; + Case.MatrixAValues = makeMatrixArithmeticPattern(Case.M, Case.K, 3, 2, 5, 2); + Case.MatrixBValues = makeMatrixArithmeticPattern(Case.K, Case.N, 1, 3, 7, 3); + if (Case.accumulates()) { + Case.AccumulatorValues = + makeMatrixArithmeticPattern(Case.M, Case.N, 2, 1, 5, 2); + Case.PublicRule = + L"Exact non-uniform ThreadGroup F16 product plus an independent F32 " + L"accumulator"; + } else if (AccumulatorType == ComponentType::F32) { + Case.PublicRule = + L"Exact non-uniform ThreadGroup F16 matrix product stored in an F32 " + L"accumulator"; + } else { + Case.PublicRule = + L"Exact non-uniform ThreadGroup F16 product with rectangular inputs"; + } + return Case; +} + +void DxilConf_SM610_LinAlg::MatMatMul_ThreadGroup_8x16x8_F16_NonUniform() { + const MatrixMultiplyCase Case = makeRectangularF16ThreadGroupMultiplyCase( + ComponentType::F16, MatrixMultiplyOperation::Multiply); + runThreadGroupMultiplyCase(D3DDevice, DxcSupport, Case, + L"MatMatMul_ThreadGroup_8x16x8_F16_NonUniform", + VerboseLogging); +} + +void DxilConf_SM610_LinAlg:: + MatMatMulAccum_ThreadGroup_8x16x8_F16_ToF32_NonUniform() { + const MatrixMultiplyCase Case = makeRectangularF16ThreadGroupMultiplyCase( + ComponentType::F32, MatrixMultiplyOperation::MultiplyAccumulate); + runThreadGroupMultiplyCase( + D3DDevice, DxcSupport, Case, + L"MatMatMulAccum_ThreadGroup_8x16x8_F16_ToF32_NonUniform", + VerboseLogging); +} + +void DxilConf_SM610_LinAlg::MatMatMul_ThreadGroup_8x8x8_I32() { + MatrixMultiplyCase Case = {}; + Case.MatrixAType = ComponentType::I32; + Case.MatrixBType = ComponentType::I32; + Case.AccumulatorType = ComponentType::I32; + Case.M = 8; + Case.K = 8; + Case.N = 8; + Case.MatrixAValues = makeMatrixArithmeticPattern(Case.M, Case.K, 3, 2, 5, 2); + Case.MatrixBValues = makeMatrixArithmeticPattern(Case.K, Case.N, 1, 3, 7, 3); + Case.PublicRule = L"Exact non-uniform ThreadGroup I32 matrix product"; + runThreadGroupMultiplyCase(D3DDevice, DxcSupport, Case, + L"MatMatMul_ThreadGroup_8x8x8_I32", + VerboseLogging); +} + static const char WaveAccumulateBUseShader[] = R"( #define USE_ACC 2