From 534928d35b089475b76cd52c444dc79e601f8869 Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub@nod-labs.com>
Date: Wed, 10 Jul 2024 13:34:54 -0400
Subject: [PATCH] [LLVMGPU] Add debug print for contraction problem size. NFC.
 (#17845)

Signed-off-by: Jakub Kuderski <jakub.kuderski@amd.com>
---
 .../compiler/Codegen/LLVMGPU/KernelConfig.cpp     | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
index 5879ac0de2e4..ed2c8e8d82a7 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
@@ -380,9 +380,9 @@ setConvolutionVectorDistributionConfig(IREE::GPU::TargetAttr target,
 }
 
 [[maybe_unused]] static void
-debugPrintContractionInfo(unsigned numLoops,
+debugPrintContractionInfo(StringRef label, unsigned numLoops,
                           linalg::ContractionDimensions contractionDims,
-                          ArrayRef<int64_t> workgroupTileSizes) {
+                          ArrayRef<int64_t> sizes) {
   ArrayRef<unsigned> dimVals[] = {contractionDims.batch, contractionDims.m,
                                   contractionDims.n, contractionDims.k};
   std::string dimSymbols(numLoops, '*');
@@ -395,8 +395,8 @@ debugPrintContractionInfo(unsigned numLoops,
   llvm::interleaveComma(dimSymbols, llvm::dbgs());
   llvm::dbgs() << "]\n";
 
-  DBGS() << "Workgroup tile sizes: [";
-  llvm::interleaveComma(workgroupTileSizes, llvm::dbgs());
+  DBGS() << label << ": [";
+  llvm::interleaveComma(sizes, llvm::dbgs());
   llvm::dbgs() << "]\n";
 }
 
@@ -419,6 +419,9 @@ setMatmulVectorDistributionConfig(IREE::GPU::TargetAttr target,
     return failure();
   }
 
+  LLVM_DEBUG(debugPrintContractionInfo("Problem size", op.getNumLoops(),
+                                       *contractionDims, bounds));
+
   // For now we are not being smart and trying to reshape dimensions to allow
   // for better usage of intrinsics, and instead are tiling all dimensions
   // except the inner most m, n, and k dimensions to 1.
@@ -577,8 +580,8 @@ setMatmulVectorDistributionConfig(IREE::GPU::TargetAttr target,
   // Follow the LLVMGPU convention of keeping all of the tile sizes in one list.
   workgroupTileSizes[kDim] = schedule->kTileCount * schedule->kSize;
 
-  LLVM_DEBUG(debugPrintContractionInfo(op.getNumLoops(), *contractionDims,
-                                       workgroupTileSizes));
+  LLVM_DEBUG(debugPrintContractionInfo("Workgroup tile sizes", op.getNumLoops(),
+                                       *contractionDims, workgroupTileSizes));
 
   TileSizesListType tileSizes;
   tileSizes.push_back(workgroupTileSizes);