Fix code generation when UseBeta is false

projects/hipblaslt/tensilelite/Tensile/KernelWriterAssembly.py

@@ -1802,16 +1802,18 @@ def loadBatchedAddress(self, kernel, Batch, tmpSgprResource: ContinuousRegister)
           module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(Batch), src1=0x8, comment="offset of global buffer address"))
           module.add(SLoadB64(dst=sgpr("AddressD", 2), base=sgpr("AddressD",2), soffset=sgpr(tmpSgpr), comment="load global buffer D address"))
 
-      endCheckLabel = Label(self.labels.getName(f"label_skip_c_buffer_deref_{Batch}"), "")
-      module.add(BranchIfZero("Beta", kernel["ProblemType"]["ComputeDataType"].toEnum(), tmpSgpr, laneSC, endCheckLabel, \
-                     kernel['WavefrontSize']))
+      # Only load C buffer address if Beta is used and potentially non-zero
+      if kernel["ProblemType"]["UseBeta"]:
+        endCheckLabel = Label(self.labels.getName(f"label_skip_c_buffer_deref_{Batch}"), "")
+        module.add(BranchIfZero("Beta", kernel["ProblemType"]["ComputeDataType"].toEnum(), tmpSgpr, laneSC, endCheckLabel, \
+                       kernel['WavefrontSize']))
 
-      for idx in kernel["ProblemType"]["IndicesBatch"]:
-        if not isPackedIndex(kernel,idx):
-          module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(Batch), src1=0x8, comment="offset of global buffer address"))
-          module.add(SLoadB64(dst=sgpr("AddressC", 2), base=sgpr("AddressC",2), soffset=sgpr(tmpSgpr), comment="load global buffer C address"))
+        for idx in kernel["ProblemType"]["IndicesBatch"]:
+          if not isPackedIndex(kernel,idx):
+            module.add(SMulI32(dst=sgpr(tmpSgpr), src0=sgpr(Batch), src1=0x8, comment="offset of global buffer address"))
+            module.add(SLoadB64(dst=sgpr("AddressC", 2), base=sgpr("AddressC",2), soffset=sgpr(tmpSgpr), comment="load global buffer C address"))
 
-      module.add(endCheckLabel)
+        module.add(endCheckLabel)
 
     #handle Batch A/B
     endCheckLabel = Label(self.labels.getName(f"label_skip_ab_buffer_deref_{Batch}"), "")
@@ -2485,9 +2487,13 @@ def calculateWG():
           moduleExternalArgs.addModuleAsFlatItems(self.externalArgLoader.loadAllKernArg(sgprStart, "KernArgAddress", load, 4))
           offset = self.externalArgLoader.getOffset() + self.states.bpr * (self.states.userArgsInfo.alphaMaxRegisterSize - self.states.numSgprAlpha)
           self.externalArgLoader.setOffset(offset)
-          moduleExternalArgs.addComment("Read Beta")
-          moduleExternalArgs.addModuleAsFlatItems(self.externalArgLoader.loadAllKernArg(self.sgprs["Beta"], "KernArgAddress", self.states.numSgprBeta))
-          offset = self.externalArgLoader.getOffset() + self.states.bpr * (self.states.userArgsInfo.betaMaxRegisterSize - self.states.numSgprBeta)
+          if kernel["ProblemType"]["UseBeta"]:
+            moduleExternalArgs.addComment("Read Beta")
+            moduleExternalArgs.addModuleAsFlatItems(self.externalArgLoader.loadAllKernArg(self.sgprs["Beta"], "KernArgAddress", self.states.numSgprBeta))
+            offset = self.externalArgLoader.getOffset() + self.states.bpr * (self.states.userArgsInfo.betaMaxRegisterSize - self.states.numSgprBeta)
+          else:
+            # Even when not using Beta, we need to skip over the Beta argument space
+            offset = self.externalArgLoader.getOffset() + self.states.bpr * self.states.userArgsInfo.betaMaxRegisterSize
           if kernel["ProblemType"]["UseScaleAB"] == "Scalar":
             sgprOffset = self.externalArgLoader.getOffset()
             for preloadScale, name in zip([self.states.preloadScaleA, self.states.preloadScaleB], ['A','B']):
@@ -13737,8 +13743,8 @@ def globalWriteElements(self, kernel, tPA, tPB, vectorWidths_2, vectorWidths_1,
         self.sgprPool.checkIn(sgprScaleA)
         self.sgprPool.checkIn(sgprScaleB)
 
-      # Update beta
-      if kernel["ProblemType"]["UseScaleCD"] and ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["StreamK"] > 0):
+      # Update beta with ScaleC (only when Beta is actually used)
+      if kernel["ProblemType"]["UseBeta"] and kernel["ProblemType"]["UseScaleCD"] and ((kernel["GlobalSplitU"] == 1 or kernel["GlobalSplitU"] == -1) or kernel["StreamK"] > 0):
         assert(kernel["ProblemType"]["ComputeDataType"].isSingle())
         newBetaVgpr = self.vgprPool.checkOut(1)
         module.add(VMovB32(dst=vgpr(newBetaVgpr), src=sgpr("Beta")))

projects/hipblaslt/tensilelite/Tensile/Tests/common/gemm/use_beta_false.yaml

@@ -0,0 +1,56 @@
+# Test for UseBeta=False functionality
+# Verifies that kernels correctly handle the case where beta=0 and tensor C is not used
+
+GlobalParameters:
+  MinimumRequiredVersion: 5.0.0
+  PrintLevel: 1
+  ForceRedoBenchmarkProblems: True
+  ForceRedoLibraryLogic: True
+  ForceRedoLibraryClient: True
+  CMakeBuildType: Release
+  EnqueuesPerSync: 1
+  SyncsPerBenchmark: 0
+  NumElementsToValidate: 128
+  Platform: 0
+  Device: 0
+  KernelTime: True
+  SleepPercent: 0
+  NumBenchmarks: 1
+  PrintSolutionRejectionReason: True
+  LibraryFormat: yaml
+  BoundsCheck: True
+
+BenchmarkProblems:
+  ########################################
+  # UseBeta=False with batched GEMM
+  ########################################
+  -
+    - # ProblemType
+      OperationType: GEMM
+      DataType: h
+      DestDataType: h
+      ComputeDataType: s
+      HighPrecisionAccumulate: True
+      TransposeA: False
+      TransposeB: True
+      UseBeta: False
+      Batched: True
+
+    - # Configuration
+      InitialSolutionParameters:
+      BenchmarkCommonParameters:
+        - KernelLanguage: ["Assembly"]
+      ForkParameters:
+        - MatrixInstruction:
+          - [16, 16, 16, 1, 1, 2, 2, 2, 2]
+        - DepthU: [16]
+        - VectorWidthA: [2]
+        - VectorWidthB: [2]
+        - GlobalSplitU: [1]
+      BenchmarkForkParameters:
+      JoinParameters:
+      BenchmarkJoinParameters:
+      BenchmarkFinalParameters:
+        - ProblemSizes:
+          - Exact: [256, 256, 1, 256]
+          - Exact: [128, 128, 1, 128]

projects/hipblaslt/tensilelite/client/include/ReferenceValidator.hpp

@@ -127,6 +127,10 @@ namespace TensileLite
         private:
             void allocateResultBuffer(size_t bytes);
 
+            bool shouldSkipNullTensor(const std::string& tensorName,
+                                      bool hasNullPointer,
+                                      bool hasZeroElements) const;
+
             std::shared_ptr<DataInitialization> m_dataInit;
             std::shared_ptr<ProblemInputs>      m_referenceInputs;
 

projects/hipblaslt/tensilelite/client/src/DataInitialization.cpp

@@ -761,12 +761,14 @@ namespace TensileLite
                                   size_t                  totalElements,
                                   hipMemcpyKind           kind)
         {
+            // First, copy the NaN-filled bad buffer to the entire destination for padding
             HIP_CHECK_EXC(
                 hipMemcpy(dst,
-                          src,
+                          bad,
                           multiplyElementSize(totalElements,
                                               DataTypeInfo::Get(descriptor.dataType()).elementSize),
                           kind));
+            // Then, copy valid data to the middle section, overwriting the NaN padding in that region
             ptrdiff_t dPadding = totalElements - descriptor.totalAllocatedElements();
             dPadding           = multiplyElementSize(dPadding, descriptor.elementBytes());
             void* dstOffset    = (void*)((uint8_t*)dst + dPadding / 2);
@@ -803,8 +805,12 @@ namespace TensileLite
                                size_t                  totalElements,
                                hipMemcpyKind           kind)
         {
-            HIP_CHECK_EXC(hipMemcpy(
-                dst, src, multiplyElementSize(totalElements, descriptor.elementBytes()), kind));
+            // Skip copy if no elements to copy or if pointers are null (e.g., when UseBeta=false, tensor C may not be allocated)
+            if(totalElements > 0 && dst != nullptr && src != nullptr)
+            {
+                HIP_CHECK_EXC(hipMemcpy(
+                    dst, src, multiplyElementSize(totalElements, descriptor.elementBytes()), kind));
+            }
             return dst;
         }
 

projects/hipblaslt/tensilelite/client/src/ReferenceValidator.cpp

@@ -35,6 +35,7 @@
 #include <Tensile/hip/HipUtils.hpp>
 
 #include <cstddef>
+#include <sstream>
 
 namespace TensileLite
 {
@@ -361,6 +362,36 @@ namespace TensileLite
             return rv;
         }
 
+        bool ReferenceValidator::shouldSkipNullTensor(const std::string& tensorName,
+                                                      bool hasNullPointer,
+                                                      bool hasZeroElements) const
+        {
+            // Skip validation if pointers are null or no data to validate
+            // Only tensor C can be null when beta is zero (UseBeta=false)
+            if(!hasNullPointer && !hasZeroElements)
+                return false;
+
+            // Tensor C is allowed to be null when beta is zero (not used)
+            bool isTensorC = (tensorName == "C");
+            bool isBetaZero = false;
+            if(m_problem != nullptr)
+            {
+                auto* gemmProblem = dynamic_cast<ContractionProblemGemm const*>(m_problem);
+                if(gemmProblem != nullptr)
+                    isBetaZero = (gemmProblem->beta() == 0.0);
+            }
+
+            if(isTensorC && isBetaZero)
+            {
+                if(Debug::Instance().printTensorInfo())
+                    std::cout << "Skipping validation for tensor C (beta=0, not used)" << std::endl;
+                return true;
+            }
+
+            // For all other cases, null pointers or no data to validate is an error
+            return false;
+        }
+
         bool ReferenceValidator::validate(ContractionProblemGemm const& problem,
                                           ContractionInputs const&      reference,
                                           ContractionInputs const&      result)
@@ -487,7 +518,26 @@ namespace TensileLite
                     std::cout << "Validating tensor " << tensor.getName() << ", cpu pointer "
                               << refPtr << ", gpu pointer " << resPtr
                               << ", size = " << result.maxElements[i] << std::endl;
-
+
+                // Check if we should skip this tensor due to null pointers or zero elements
+                bool hasNullPointer = (resPtr == nullptr || refPtr == nullptr);
+                bool hasZeroElements = (result.maxElements[i] == 0);
+
+                if(shouldSkipNullTensor(tensor.getName(), hasNullPointer, hasZeroElements))
+                {
+                    continue;
+                }
+
+                // If we reach here with null pointers or zero elements, it's an error
+                if(hasNullPointer || hasZeroElements)
+                {
+                    std::stringstream ss;
+                    ss << "Unexpected null pointer or zero elements for tensor " << tensor.getName()
+                       << " (resPtr=" << resPtr << ", refPtr=" << refPtr
+                       << ", maxElements=" << result.maxElements[i] << ")";
+                    throw std::runtime_error(ss.str());
+                }
+
                 rv &= checkResults(
                     tensor, refPtr, resPtr, result.maxElements[i], result.gpu, validationStride, threshold);
             }
@@ -496,13 +546,15 @@ namespace TensileLite
 
         void ReferenceValidator::allocateResultBuffer(size_t bytes)
         {
-            if(m_cpuResultBufferSize == bytes)
+            // Only skip reallocation if size matches AND buffer is valid
+            if(m_cpuResultBufferSize == bytes && m_cpuResultBuffer.get() != nullptr)
                 return;
+
             m_cpuResultBuffer.reset();
 
             uint8_t* buffer;
-            HIP_CHECK_EXC(hipHostMalloc(&buffer, bytes, 0));
-            m_cpuResultBuffer.reset(buffer, hipFree);
+            HIP_CHECK_EXC(hipHostMalloc((void**)&buffer, bytes, 0));
+            m_cpuResultBuffer.reset(buffer, [](uint8_t* p) { (void)hipFree(p); });
             m_cpuResultBufferSize = bytes;
         }
 
@@ -685,26 +737,52 @@ namespace TensileLite
             size_t elementsAfterData    = 0;
 
             BoundsCheckMode boundsCheck = m_dataInit->getCurBoundsCheck();
+            // For NaN bounds checking, copy the full padded buffer from GPU for all tensors
             if(boundsCheck == BoundsCheckMode::NaN)
                 elementsToCopy = maxElement;
             size_t bytesToCopy = elementsToCopy * sizeof(ValidType);
 
-            if(m_cpuResultBufferSize < bytesToCopy)
-                allocateResultBuffer(bytesToCopy);
+            // Check if we should skip this tensor due to null pointers or no data
+            bool hasNullPointer = (result == nullptr || reference == nullptr);
+            bool hasZeroElements = (bytesToCopy == 0 || maxElement == 0);
 
-            auto copykind = isgpu ? hipMemcpyDeviceToHost : hipMemcpyHostToHost;
+            if(shouldSkipNullTensor(tensor.getName(), hasNullPointer, hasZeroElements))
+            {
+                return true;
+            }
 
+            // If we reach here with null pointers or no data, it's an error
+            if(hasNullPointer || hasZeroElements)
             {
-                ScopedTimer timer("validate_gpu_readback");
-                HIP_CHECK_EXC(hipMemcpy(m_cpuResultBuffer.get(), result, bytesToCopy, copykind));
+                std::stringstream ss;
+                ss << "Unexpected null pointer or no data for tensor " << tensor.getName()
+                   << " (result=" << result << ", reference=" << reference
+                   << ", bytesToCopy=" << bytesToCopy << ", maxElement=" << maxElement << ")";
+                throw std::runtime_error(ss.str());
             }
 
+            if(m_cpuResultBufferSize < bytesToCopy || m_cpuResultBuffer.get() == nullptr)
+                allocateResultBuffer(bytesToCopy);
+
+            auto copykind = isgpu ? hipMemcpyDeviceToHost : hipMemcpyHostToHost;
+
+            // For NaN bounds checking, the result pointer points to valid data (middle of buffer)
+            // We need to adjust it back to buffer start to copy the NaN padding
+            void const* copySource = result;
             if(boundsCheck == BoundsCheckMode::NaN)
             {
                 ptrdiff_t bPadding = maxElement - tensor.totalAllocatedElements();
                 elementsBeforeData = bPadding / 2;
                 elementsAfterData
                     = elementsToCopy - (tensor.totalAllocatedElements() + elementsBeforeData);
+                // Adjust pointer back to buffer start (before the padding)
+                copySource = (uint8_t const*)result
+                             - elementsBeforeData * sizeof(ValidType);
+            }
+
+            {
+                ScopedTimer timer("validate_gpu_readback");
+                HIP_CHECK_EXC(hipMemcpy(m_cpuResultBuffer.get(), copySource, bytesToCopy, copykind));
             }
             // If there was extra data allocated before the tensor to do bounds
             // checking, resultBuffer is the whole allocation, while resultData