[Mlir-commits] [mlir] b96d069 - [NVGPU] Add debug in nvgpu (nfc)

[Guray Ozen]

Author: Guray Ozen
Date: 2023-09-01T16:40:40+02:00
New Revision: b96d0693244be768ad5feec537cf08f6e8507ee0

URL: https://github.com/llvm/llvm-project/commit/b96d0693244be768ad5feec537cf08f6e8507ee0
DIFF: https://github.com/llvm/llvm-project/commit/b96d0693244be768ad5feec537cf08f6e8507ee0.diff

LOG: [NVGPU] Add debug in nvgpu (nfc)

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D159343

Added: 
    

Modified: 
    mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
index 9073439237463e..b045089244ff1a 100644
--- a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
+++ b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
@@ -20,8 +20,13 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Pass/Pass.h"
+#include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
+#define DEBUG_TYPE "nvgpu-to-nvvm"
+#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
+#define DBGSE() (llvm::dbgs())
+
 namespace mlir {
 #define GEN_PASS_DEF_CONVERTNVGPUTONVVMPASS
 #include "mlir/Conversion/Passes.h.inc"
@@ -980,9 +985,14 @@ struct NVGPUGenerateGmmaDescriptorLowering
     };
 
     int ex4LSB = 4;
-    Value strideDim = makeConst((layout << 3) >> ex4LSB);
     int64_t sizeN = op.getTensorMap().getType().getTensor().getDimSize(0);
-    Value leadDim = makeConst((sizeN * layout) >> ex4LSB);
+    uint64_t strideDimVal = (layout << 3) >> ex4LSB;
+    uint64_t leadDimVal = (sizeN * layout) >> ex4LSB;
+    uint64_t offsetVal = 0;
+
+    Value strideDim = makeConst(strideDimVal);
+    Value leadDim = makeConst(leadDimVal);
+
     Value baseAddr = getStridedElementPtr(
         op->getLoc(), cast<MemRefType>(op.getTensor().getType()),
         adaptor.getTensor(), {}, rewriter);
@@ -996,7 +1006,7 @@ struct NVGPUGenerateGmmaDescriptorLowering
     // // [62,64)  swizzle type
     dsc = insertBit(dsc, makeConst(swizzle), startSwizzleBit);
     // // [49,52)  base_offset
-    dsc = insertBit(dsc, makeConst(0), startOffsetBit);
+    dsc = insertBit(dsc, makeConst(offsetVal), startOffsetBit);
     // // [32,46)  stride
     dsc = insertBit(dsc, strideDim, startStrideBit);
     // // [16,30)  leading dimension
@@ -1004,6 +1014,14 @@ struct NVGPUGenerateGmmaDescriptorLowering
     // // [0,14)   start_address
     dsc = insertBit(dsc, basePtr14bit, startBaseAddrBit);
 
+    LLVM_DEBUG(DBGS() << "Generating wgmma.descriptor: "
+                      << "leading_off:" << leadDimVal << "\t"
+                      << "stride_off :" << strideDimVal << "\t"
+                      << "base_offset:" << offsetVal << "\t"
+                      << "layout_type:" << swizzle << " ("
+                      << nvgpu::stringifyTensorMapSwizzleKind(swizzleKind)
+                      << ")\n start_addr :  " << baseAddr << "\n");
+
     rewriter.replaceOp(op, dsc);
     return success();
   }