/******************************************************************************* * * MIT License * * Copyright (C) 2022-2026 Advanced Micro Devices, Inc. All rights reserved. * * Permission is hereby granted, free of charge, to any person obtaining a copy * of this software and associated documentation files (the "Software"), to deal * in the Software without restriction, including without limitation the rights * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell * copies of the Software, and to permit persons to whom the Software is * furnished to do so, subject to the following conditions: * * The above copyright notice and this permission notice shall be included in * all copies or substantial portions of the Software. * * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE * SOFTWARE. * *******************************************************************************/ #pragma once #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include "origami/origami.hpp" #include "origami/streamk.hpp" #include #define TENSILE_COMMON_KERNEL_ARGS_SIZE 16 namespace TensileLite { template struct DeviceUserArguments { uint32_t m; uint32_t n; uint32_t batch; uint32_t k; void* d; void* c; void* a; void* b; uint32_t strideD1; uint32_t strideD2; uint32_t strideC1; uint32_t strideC2; uint32_t strideA1; uint32_t strideA2; uint32_t strideB1; uint32_t strideB2; int8_t alpha[16]; int8_t beta[16]; void* scaleA; void* scaleB; void* scaleC; void* scaleD; void* scaleAlphaVec; void* bias; int biasType; uint32_t reserved; void* e; uint32_t strideE1; uint32_t strideE2; TAct act0; TAct act1; int activationType; } __attribute__((packed)); struct PerfModel { double clock = std::numeric_limits::quiet_NaN(); double memClock = std::numeric_limits::quiet_NaN(); double peakGFlops = std::numeric_limits::quiet_NaN(); double memBandwidthMBps = std::numeric_limits::quiet_NaN(); double l2ReadBwMul = std::numeric_limits::quiet_NaN(); double gFlops = std::numeric_limits::quiet_NaN(); double readEff = 0.0; double l2ReadHitRate = 0.0; double l2WriteHitRate = 0.0; int CUs = 0; }; extern TENSILELITEHOST_EXPORT PerfModel perf; struct BufferLoadCheckPacket { size_t shiftPtrElemA; size_t shiftPtrElemB; size_t depthUorMT0; size_t depthUorMT1; }; struct SizeMapping { size_t waveNum; dim3 clusterDim{1, 1, 1}; dim3 workGroupSize; dim3 threadTile; dim3 macroTile; std::array matrixInstruction; size_t grvwA = 1; size_t grvwB = 1; size_t gwvwC = 1; size_t gwvwD = 1; size_t staggerU = 0; size_t staggerUMapping = 0; size_t depthU = 0; size_t globalSplitUPGR = 0; int16_t globalSplitU = 0; size_t staggerStrideShift = 0; int workGroupMapping = 0; size_t packBatchDims = 0; int packSummationDims = 0; int magicDivAlg = 1; int streamK = 0; int streamKForceDPOnly = 0; int streamKAtomic = 0; int prefetchAcrossPersistent = 0; int persistentKernel = 0; bool persistentKernelAlongBatch = false; bool sourceKernel = false; int globalAccumulation = 0; int adaptiveGemmGSUA = 0; size_t workspaceSizePerElemC = 0; size_t workspaceSizePerElemBias = 0; bool activationFused = true; std::string customKernelName; int workGroupMappingXCC = 0; int workGroupMappingXCCGroup = 0; bool globalSplitUCoalesced = false; bool globalSplitUWorkGroupMappingRoundRobin = false; int CUOccupancy = 0; int PrefetchGlobalRead = 2; int MathClocksUnrolledLoop = 0; size_t synchronizerSizePerWG = 0; int nonTemporalA = 0; int nonTemporalB = 0; int adaptiveGemmNTAB = 0; int customMainLoopScheduling = 0; // Whether the kernel uses the subtile implementation (UseSubtileImpl). // Plumbed into the Origami config so heuristics can reason about subtile // kernels (e.g. rejecting them for small K). bool useSubtileImpl = false; int NonTemporalD = 0; int WaveSeparateGlobalReadA = 0; int WaveSeparateGlobalReadB = 0; int UnrollLoopSwapGlobalReadOrder = 0; bool DirectToVgprA = false; bool DirectToVgprB = false; int NumLoadsCoalescedA = 0; int NumLoadsCoalescedB = 0; int VectorWidthA = 1; int VectorWidthB = 1; int LocalSplitU = 1; bool DirectToLdsA = false; bool DirectToLdsB = false; int expertSchedulingMode = 0; std::array waveGroup; }; struct StreamKSettings { origami::reduction_t reduction = origami::reduction_t::tree; size_t grid = 0; // StreamK=5 tri-state (0=OFF default/SK3, 1=ON/SK4, 2=AUTO); see // hipblasLtStreamKTileSchedulingMode_t. Ignored when streamK != 5. int streamKTileSchedulingMode = 0; int smCountTarget = 0; // 0 = use all device CUs; >0 engages origami heuristic when mode is OFF }; struct GSUSettings { size_t globalAccumulation = 0; }; /** * Represents a single kernel or set of kernels that can perform a single * tensor contraction. * * Can generate `KernelInvocation` objects to solve a particular problem * given a set of `ContractionInputs`. */ class TENSILELITEHOST_EXPORT ContractionSolution : public Solution { public: using Problem = ContractionProblemGemm; using Inputs = ContractionInputs; using GroupedInputs = ContractionGroupedInputs; using WGMParamsCache = CacheMap, Problem>; using StaggerUParamsCache = CacheMap, Problem>; /** * Indicate a solution's matching type */ enum class MatchingTag { Equal, // EqualityMatching Range, // RangeMatching FreeSize, // FreeSizeMatching GridBased, // GridBasedMatching Prediction, // PredictionMatching Experimental, // ExperimentalStreamK or ExperimentalMLP Others, // Default }; static std::string Type() { return "Contraction"; } virtual std::string type() const { return Type(); } virtual std::string KernelName() const { return kernelName; } virtual std::string name() const { return solutionName; } virtual std::string description() const { return kernelName; } virtual bool isFallbackForHW(Hardware const&) const; bool isStreamK() const { return sizeMapping.streamK > 0; } /** * @brief Returns the string representation of the solution's matching type. * * This tag is used to identify or categorize the solution for matching purposes. * @return A string representing the matching type of the solution. */ virtual std::string matchingTag() const; //! Estimates based on problem size, solution tile, and machine hardware //! charz: struct StaticPerformanceModel { size_t memReadBytesA = 0.0; //! Estimated memory reads A size_t memReadBytesB = 0.0; //! Estimated memory reads B size_t memReadBytesC = 0.0; //! Estimated memory reads C size_t memWriteBytesD = 0.0; //! Estimated memory writes D size_t memReadBytes = 0.0; size_t memGlobalReads = 0; size_t memGlobalWrites = 0; }; struct Granularities { double numTiles0 = 0.0; //! number of tiles in 0 dimension double numTiles1 = 0.0; //! number of tiles in 1 dimension double totalTiles = 0.0; double tilesPerCu = 0.0; //! Granularity is measured 0..1 with 1.0 meaning no granularity loss double tile0Granularity = 0.0; // loss due to tile0 double tile1Granularity = 0.0; double cuGranularity = 0.0; double waveGranularity = 0.0; double totalGranularity = 0.0; double totalTileAwareGranularity = 0.0; double natCuGranularity = 0.0; double natTilesPerCu = 0.0; double suTilesPerCu = 0.0; double suCuGranularity = 0.0; double waves = 0.0; double suWavesPerSimdx2 = 0.0; double suWaveGranularity = 0.0; int CUs = 0; double MT0; double MT1; double GSU; double LSU; }; struct ProjectedPerformance { Granularities granularities; double speedGFlops = 0.0; //! final gflops projection int CUs = 0; StaticPerformanceModel staticModel; }; // Result of host-side AdaptiveGemmNTAB dispatch. // nta/ntb are either 0 (cached) or 4 (non-temporal). // Packed into internalArg0 bits 12/13 when InternalArgsSupport.version >= 3. struct AdaptiveGemmNTAB { uint32_t nta = 0; uint32_t ntb = 0; }; bool checkInternalArgumentsSupport(ContractionProblem const& problem, std::ostream& stream, bool debug = false) const; /** * Calculate required workspace size. */ size_t requiredWorkspaceSize(Problem const& problem, Hardware const& hardware) const; size_t requiredWorkspaceSizeGsu(Problem const& problem, Hardware const& hardware, size_t gsu) const; size_t requiredWorkspaceSizeGroupedGemm(std::vector const& problems, Hardware const& hardware) const; size_t requiredHostSizeGroupedGemmSingle(Problem const& problem, Hardware const& hardware) const; size_t requiredSynchronizerSize(Problem const& problem, Hardware const& hardware) const; void calculateGrid(dim3& workGroupSize, dim3& numWorkGroups, ContractionSolution::Problem const& problem) const; origami::reduction_t getSKReduction(Problem const& problem, Hardware const& hardware) const; size_t getSKGrid(Problem const& problem, Hardware const& hardware, size_t tiles, origami::reduction_t reductionStrat) const; // Resolve the effective StreamK=5 hybrid sub-mode for a launch: returns // true for the dynamic (SK4) path, false for the static (SK3) path. // Precedence (highest first): the TENSILE_STREAMK5_FORCE_MODE debug env // override (0=force static, 1=force dynamic), then the problem tri-state // streamKTileSchedulingMode (0=OFF/static unless smCountTarget()>0, // 1=ON/dynamic), then AUTO (2) via the origami hybrid-mode heuristic. // Only meaningful when // sizeMapping.streamK == 5. This is the single source of truth shared by // grid sizing (getSKGrid) and kernel-arg packing (generateSingleCall) so // the launch grid and the packed args can never disagree. bool streamK5EffectiveDynamic(Problem const& problem, Hardware const& hardware) const; // Selection-time predicate for the StreamK dynamic-queue / work-stealing // path. The SK4 and dynamic sub-path of SK5 kernels hardcode a // power-of-two per-XCD queue count and mask indices with (Q-1); that fast // masking is only valid when the device exposes a power-of-two number of // XCDs. Returns false (and warns once) when this solution would take the // dynamic-queue path but the hardware's NUM_XCD is not a power of two // (e.g. MI300A = 6), so the solution is EXCLUDED from selection rather // than silently degraded to tree reduction. All other solutions return // true. Wired into softwarePredicate() (SolutionLibrary.hpp). bool streamKDynamicQueueSupported(Problem const& problem, Hardware const& hardware) const; size_t partialTileSize(size_t skGrid) const; static float computeGranularity(float x); Granularities computeGranularities(Hardware const& hardware, double M, double N, double K, double NumBatches, uint32_t autoGsuVal) const; StaticPerformanceModel staticPerformanceModel(double M, double N, double K, double NumBatches, double MT0, double MT1, double NumCUs, double totalGranularity, int globalSplitU) const; /** * Calculate the projected performance based on granularity loss. */ ProjectedPerformance projectedPerformance(Problem const& problem, Hardware const& hardware) const; /** * Generate a set of kernel calls to solve a particular problem. */ virtual std::vector solve(ContractionProblem const& problem, ProblemInputs const& inputs, Hardware const& hardware, void* hipHostMemory, size_t hipHostMemorySize, hipStream_t stream) const; virtual std::vector solve(Problem const& problem, Inputs const& inputs, Hardware const& hardware) const; virtual std::vector solveGroupedGemm(std::vector const& problems, GroupedInputs const& inputs, Hardware const& hardware, void* hipHostMemory, size_t hipHostMemorySize, hipStream_t stream) const; // The problems and inputs are passed by device memory virtual std::vector solveGroupedGemmGPU(std::vector const& problems, GroupedInputs const& inputs, Hardware const& hardware, const void* dUA, const void* workspace, hipStream_t stream) const; // For Tensile debugging, will allocate and initialize DeviceUserArguments with the problems and inputs. virtual std::vector solveTensileGPU(ContractionProblem const& problem, ProblemInputs const& inputs, Hardware const& hardware, void** dUA, void** dUAHost, void* hipHostMemory, size_t hipHostMemorySize, hipStream_t stream) const; // For Tensile debugging, will allocate and initialize DeviceUserArguments with the problems and inputs. virtual std::vector solveTensileGroupedGemmGPU(std::vector const& problems, GroupedInputs const& inputs, Hardware const& hardware, void** dUA, void** dUAHost, void* hipHostMemory, size_t hipHostMemorySize, hipStream_t stream) const; virtual void relaseDeviceUserArgs(void* dUA, void* dUAHost); template void singleCallArgs(Problem const& problem, ContractionInputs const& inputs, uint32_t const& workspaceOffsetInByte, Hardware const* hardware, dim3 const& problemNumGroupTiles, dim3 const& numWorkGroups, KA& args, StreamKSettings const& sk) const; // Common kernel related arguments (e.g. gemm_count, arg type, MT, GSU...) template void kernelArgs(uint32_t gemmCount, uint32_t argType, KA& args, uint32_t numWorkGroups, Hardware const* hardware, const ContractionProblemParameters& param, int32_t autoWGM, size_t autoWGMXCC, size_t autoWGMXCCCHUNK, size_t autoStaggerUMapping, size_t autoStaggerU, size_t autoStaggerUStrideShift, uint32_t autoGsuVal, AdaptiveGemmNTAB ntab) const; template inline void calculateSingleCallWorkGroupItems(std::vector const& problems, const TensileLite::dim3& workGroupSize, TensileLite::dim3& numWorkGroups, TensileLite::dim3& numWorkItems, KA& h_args, uint32_t gsu) const; template KernelInvocation generateSingleCall(Problem const& problem, ContractionInputs const& inputs, Hardware const& hardware, StreamKSettings const& sk, GSUSettings const& gsuSettings) const; template KernelInvocation generateSingleCallGroupedGemm(std::vector const& problems, GroupedInputs const& inputs, Hardware const& hardware, KA& h_args, void const* userArgs = nullptr) const; template KernelInvocation generateBetaOnlyCall(Problem const& problem, ContractionInputs const& inputs) const; template KernelInvocation generateBetaOnlyCallGroupedGemm(std::vector const& problems, GroupedInputs const& inputs) const; std::string betaOnlyKernelName(Problem const& problem) const; template void outputConversionCallArgs(Problem const& problem, ContractionInputs const& inputs, uint32_t const& workspaceOffsetInByte, KA& args, StreamKSettings const& sk, uint32_t autoGsuVal, uint32_t additionalPaddingPerBatchGeneralBatch=0) const; template inline void calculateConversionCallWorkGroupItems( std::vector const& problems, size_t& vw, const TensileLite::dim3& workGroupSize, TensileLite::dim3& numWorkGroups, TensileLite::dim3& numWorkItems, KA& args) const; template KernelInvocation generateOutputConversionCall(Problem const& problem, ContractionInputs const& inputs, StreamKSettings const& sk, uint32_t autoGsuVal) const; template KernelInvocation generateOutputConversionCallGroupedGemm(std::vector const& problems, GroupedInputs const& inputs, Hardware const& hardware, KA& h_args) const; template KernelInvocation updateUserArgsOutputConversionCallGroupedGemm( std::vector const& problems, const void* userArgs, const void* workspace) const; std::string outputConversionKernelName(Problem const& problem, ContractionInputs const& inputs, size_t vw, size_t gsu) const; template KernelInvocation generateReductionCall(Problem const& problem, ContractionInputs const& inputs) const; std::string outputReductionKernelName(Problem const& problem, ContractionInputs const& inputs, size_t mt0, size_t mt1, size_t vw) const; struct InternalArgsSupport { int version = 0; bool gsu = true; bool wgm = true; bool staggerU = true; bool useUniversalArgs = true; bool useSFC = false; }; struct ProblemType { std::string operationIdentifier; bool transA = false; bool transB = false; rocisa::DataType aType = rocisa::DataType::Float; rocisa::DataType bType = rocisa::DataType::Float; rocisa::DataType cType = rocisa::DataType::Float; rocisa::DataType dType = rocisa::DataType::Float; rocisa::DataType eType = rocisa::DataType::Float; rocisa::DataType computeInputTypeA = rocisa::DataType::Float; rocisa::DataType computeInputTypeB = rocisa::DataType::Float; rocisa::DataType computeType = rocisa::DataType::Float; rocisa::DataType f32XdlMathOp = rocisa::DataType::Float; rocisa::DataType activationComputeDataType = rocisa::DataType::Float; bool highPrecisionAccumulate = false; bool useBeta = true; bool useGradient = false; int useBias = 0; bool useE = false; bool useGateResidual = false; std::string useScaleAB = ""; bool useScaleCD = false; int useScaleAlphaVec = 0; bool useInitialStridesAB = false; bool useInitialStridesCD = false; bool stridedBatched = true; bool outputAmaxD = false; bool groupedGemm = false; ActivationType activationType = ActivationType::None; int activationArgLength = 0; bool activationNoGuard = false; std::vector biasSrcWhiteList; std::vector biasDataTypeWhiteList; std::vector gateResidualDataTypeWhiteList; int sparse = 0; bool stochasticRounding = false; bool supportDeviceUserArguments = false; bool swizzleTensorA = false; bool swizzleTensorB = false; int metadataLayout = 0; int mxBlockA = 0; int mxBlockB = 0; rocisa::DataType mxTypeA = rocisa::DataType::E8; rocisa::DataType mxTypeB = rocisa::DataType::E8; // In-device MX scale layout expected by the kernel. Mirrors the // MXScaleFormat solution parameter (see Tensile/Common/ValidParameters.py). // Encoded as a small int so it round-trips through msgpack and YAML // logic files without an explicit enum schema: // 0 = NoSwizzle (default; canonical row/column layout) // 1 = HostPreSwizzle (gfx950 subtile host-preswizzled layout) // 2 = InMemorySwizzle (gfx1250 TDM-populated swizzled layout) // The host (DataInitialization) consults this to decide whether to // apply the K-dimension swizzle on the MX scale tensor before upload. int mxScaleFormat = 0; }; struct LinearModel { double slope = 1.0; double intercept = 0.0; double max = 1000.0; }; int index = 0; std::string kernelName; std::string solutionName; ThreadSafeValue codeObjectFilename; bool debugKernel = false; bool kernelArgsLog = false; mutable int isFallbackCUSol = -1; // -1:unset, 0:false, 1:true mutable WGMParamsCache wgmParamsCache = WGMParamsCache(std::make_tuple(INT32_MAX, SIZE_MAX, SIZE_MAX)); mutable StaggerUParamsCache staggerUParamsCache = StaggerUParamsCache(std::make_tuple(SIZE_MAX, SIZE_MAX, SIZE_MAX)); std::shared_ptr> taskPredicate = std::make_shared>(); std::shared_ptr> problemPredicate = std::make_shared>(); std::shared_ptr> hardwarePredicate = std::make_shared>(); SizeMapping sizeMapping; InternalArgsSupport internalArgsSupport; ProblemType problemType; // This will be calculated when getSolution is called size_t requiredHostWorkspaceSizePerProblem = static_cast(-1); /// Debugging purposes. Shouldn't contain any vital information that isn't /// somewhere else. int32_t libraryLogicIndex = -1; std::map ideals; LinearModel linearModel; MatchingTag tag{MatchingTag::Others}; uint32_t magicNumberAlg1(uint32_t x, uint32_t* magicShift) const; uint32_t magicNumberAlg2(uint32_t x, uint32_t* magicShift) const; uint32_t magicNumber(int magicDivAlg, uint32_t x, uint32_t* magicShift) const; uint32_t smallMagicNumber(uint32_t x) const; std::tuple calculateAutoWGM(Problem const& problem, Hardware const* hardware, uint32_t skgrid) const; std::tuple calculateAutoStaggerU(Problem const& problem, Hardware const* hardware, uint32_t skgrid, int32_t autoWGM) const; uint32_t calculateAutoGSU(Problem const& problem, Hardware const* hardware) const; double calculateDimensionM(Problem const& problem) const; double calculateDimensionN(Problem const& problem) const; double calculateNumBatches(Problem const& problem) const; SizeMapping getSizeMapping(void) const {return sizeMapping;}; origami::data_type_t getOrigamiDatatype(Problem const& problem) const; AdaptiveGemmNTAB calculateAdaptiveGemmNTAB(Problem const& problem, Hardware const* hardware) const; }; template TENSILELITEHOST_EXPORT void setDeviceUserArgs(std::vector const& problems, ContractionSolution::GroupedInputs const& inputs, DeviceUserArguments* args); TENSILELITEHOST_EXPORT std::ostream& operator<<(std::ostream& stream, ContractionSolution::StaticPerformanceModel const& spm); TENSILELITEHOST_EXPORT std::ostream& operator<<(std::ostream& stream, ContractionSolution::ProjectedPerformance const& spm); TENSILELITEHOST_EXPORT std::ostream& operator<<(std::ostream& stream, BufferLoadCheckPacket const& st); } // namespace TensileLite