[BACKEND] Turn on bit 46 for descriptors in mmav5 (triton-lang#8032)

Following
https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-shared-memory-descriptor

Author: lezcano

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAHelpers.h

@@ -72,6 +72,24 @@ class DotOpMmaV3SmemLoader : public DotOpMmaMemLoader {
   Value descriptor;
 };
 
+// Helper class to load shared memory slices following MMAv5 layout.
+class DotOpMmaV5SmemLoader : public DotOpMmaV3SmemLoader {
+public:
+  using DotOpMmaV3SmemLoader::DotOpMmaV3SmemLoader;
+
+  // Return a descriptor pointing to the shared memory slice at coordinates (a,
+  // b), with bit 46 set.
+  Value smemLoad(int a, int b, ConversionPatternRewriter &rewriter,
+                 Location loc) const {
+    auto tb = TritonLLVMOpBuilder(loc, rewriter);
+    Value desc = DotOpMmaV3SmemLoader::smemLoad(a, b, rewriter, loc);
+    // Set bit 46 as per
+    // https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-shared-memory-descriptor
+    Value mask = tb.int_val(64, 1ULL << 46);
+    return tb.or_(desc, mask, /*disjoint*/ true);
+  }
+};
+
 // Helper class to load tensor memory following MMAv5 layout.
 class DotOpMmaV5TmemLoader : public DotOpMmaMemLoader {
 public:

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp

@@ -454,13 +454,13 @@ void convertDotImpl(const LLVMTypeConverter &typeConverter,
                                                      interleaved, transA);
   } else {
     auto allocShapeA = getAllocShape(aTensorTy, 1);
-    aLoader = std::make_unique<DotOpMmaV3SmemLoader>(
+    aLoader = std::make_unique<DotOpMmaV5SmemLoader>(
         a, baseA, shapeA, allocShapeA, zero, 1, transA, aOperandShape,
         op.numBitsPerElementA, rewriter, loc);
   }
 
   auto allocShapeB = getAllocShape(bTensorTy, 0);
-  DotOpMmaV3SmemLoader bLoader = DotOpMmaV3SmemLoader(
+  DotOpMmaV5SmemLoader bLoader = DotOpMmaV5SmemLoader(
       b, baseB, shapeB, allocShapeB, zero, 1, transB, {mmaSizeN, mmaSizeK},
       op.numBitsPerElementB, rewriter, loc);
 

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TensorMemoryToLLVM.cpp

@@ -1016,7 +1016,7 @@ static void copySharedToTmem(ConversionPatternRewriter &rewriter, Location loc,
   auto createCopy = [&](int repM, int repN) {
     Value zero = b.i32_val(0);
     SmallVector<int64_t> shape(op.getSrc().getType().getShape());
-    DotOpMmaV3SmemLoader smemLoader = DotOpMmaV3SmemLoader(
+    DotOpMmaV5SmemLoader smemLoader = DotOpMmaV5SmemLoader(
         op.getSrc(), baseSrc, shape, op.getSrc().getType().getAllocShape(),
         zero, 1, /*trans=*/false, {128, 8},
         op.getSrc().getType().getElementType().getIntOrFloatBitWidth(),