commit 534928d35b089475b76cd52c444dc79e601f8869
author Jakub Kuderski <jakub@nod-labs.com> Wed Jul 10 13:34:54 2024 -0400
committer GitHub <noreply@github.com> Wed Jul 10 13:34:54 2024 -0400
tree 9023bdacf10f88756b4321e91ea3e84a420c4118
parent 78c0051892a87f183611abc786f72c0da0e44668

[LLVMGPU] Add debug print for contraction problem size. NFC. (#17845)

Signed-off-by: Jakub Kuderski <jakub.kuderski@amd.com>

compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp

diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
index 5879ac0..ed2c8e8 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
@@ -380,9 +380,9 @@
 }
 
 [[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 @@
   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 @@
     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 @@
   // 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);