[MLIR][NVGPU] Introduction of wgmma.generate.descriptor Op

[guraypp]

Summary

This work introduces a new Op, wgmma.generate.descriptor, designed to create a wgmma descriptor for inputs of matrix multiply and accumulate operations using wgmma.mma_async PTX instruction.

The descriptor format specifications can be found in the following link:
https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#asynchronous-warpgroup-level-matrix-shared-memory-layout-matrix-descriptor

It's important to note that this op is in its initial phase, and it does come with certain limitations. It only supports 128b swizzling and does not incorporate interleaving. In the future, different calculations will be addressed in separate works, expanding the capabilities of the op.

[qcolombet]

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td

646

Maybe add: Where Mod is the swizzling mode.

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

939

Stupid question, do we have to unwrap the descriptor, as opposed to using it directly?

978

That's not enough to get the start address, is it?

The first three fields (start address, leading dim, and offset) are stored as matrix-descriptor-encode(x) = (x & 0x3FFFF) >> 0x4, if I'm not mistaken. (Though the spec is weird because I feel we would lose some information in the encoding process for the offset and leading dim at least.)

Anyhow, this leads back to my other question, should we even unwrap the descriptor.

981

Could you introduce constants for the various shifts and sizes?
E.g., StartAddrSizeInBits = ...
StartAddrBitStartPos = ...
And use that in the shifts and masks etc.

mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir

654

I found it strange to see a 0xf16.
Could we use f16 instead?

Also I would expect generally speaking we would have something like memref<?xf16> for this kind of global variable, wouldn't we?

[qcolombet]

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

939

Scratch that, I'm stupid, we're wrapping the value in here, not unwrapping it, so we need to create that thing.
I thought there was a dedicated nvidia instruction that creates the matrix descriptor, but no :).

[guraypp]

address comments

[guraypp]

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

939

We had chat offline. Here we create descriptor from scratch, so we need to fill the bits.

978

Good catch. I added this to other fields as well. Thanks

mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir

654

Very good question.

0xf16,3 is needed for dynamic shared memory, NVPTX backend of LLVM implemented in that way.

memref<?xf16,3> IR does not get verified as a global object.

One can think of using memref<8192xf16,3>. However, there is CUDA limitation here. Sized shared memory is generated as static shared memory, and its limit is 48k whereas dynamic shared memory (0xf16) has larger.

I had chat about that with @ftynse. I am planning to model dynamic shared memory in MLIR.

[guraypp]

calculate stride at compile-time

[guraypp]

calculate leading dim also at compile time, and exclude 4 LSB

[guraypp]

use const for LSB bit

[guraypp]

fix the test

[qcolombet]

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

940

Stick to swizzle mode or say "layout mode" for the swizzle bits, because it can be difficult to connect both information when the comment explaining the struct of the bits says swizzle and the variable says layoutBit.

944

That should be 0 I believe.

982

I would suggest a higher level API (to hide the left/right shift used for masking) of the from:

val = insertBits(dst, field, start, size)

And chain that together, e.g.:

addressWo4LSB = makeShl(start_address, 4)
desc = insertBits(zero, addressWo4LSB, startAddressPos /*i.e., 0*/, startAddressSize /*i.e., 14*/)
strideDimWo4LSB = makeShl(strideDim, 4)
desc = insertBits(desc, strideDimWo4LSB, startStridePos, strideSize)
...

Where insertBits would produce something like:

insertBits(dst, field, start, size):
   mask = (1 << size) - 1 // computed as a constant at compile time
   masked_field = field & mask
   res = dst | (masked_field << start)
   return res

Note: at this point, you can simply implement exclude4LSB with makeShl(..., 4). (I.e., keep that function, the name helps the understanding, but the implementation becomes simpler.)

[qcolombet]

mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir

654

That's fine for now.
Maybe just add a comment saying that's how dynamic shared memory is currently represented.

[guraypp]

address comments

[qcolombet]

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

963

We'll also need the size of the field (e.g., 14 for BaseAddr) to mask the bits appropriately.
Technically we may not need that because our values should always be in range (e.g., the masking should be a noop), but I'd rather we emit the proper code sequence and have it being optimized away.

Alternatively, we could emit llvm.assume(val < ((1 << size) - 1)) if you believe the masking is overkill.

[guraypp]

simplify the pattern, return ssa values

[guraypp]

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

963

For the baseAddr, I do following below. Is it not enough?

Value basePtr14bit = shiftRight(shiftLeft(basePtr, 46), 50);

[qcolombet]

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

963

Ah right.
It is enough, though it is not super readable (e.g., I missed it :)).

I would have hidden the left/right shift directly in insertBits by passing a size directly. Anyhow, aside from BaseAddr I see that everything is known at compile time, so let's go with what you have here. The systematic masking is not necessary / overkill at this point.

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td

let arguments = (ins AnyUnrankedMemRef:$tensor,
                        Variadic<Index>:$boxDimensions);
   let results = (outs NVGPU_TensorMapDescriptor:$tensorMap);
   let assemblyFormat = [{
          $tensor `box` `[` $boxDimensions `]` attr-dict `:` type($tensor) `->` type($tensorMap)
   }];
   let hasVerifier = 1;
 }
 
+def NVGPU_GenerateGmmaDescriptorOp : NVGPU_Op<"wgmma.generate.descriptor", []> {
+  let summary = "Generate a wgmma matrix descriptor";
+  let description = [{
+  This Op builds a wgmma descriptor that is used by wgmma matrix multiply 
+  and accumulate.
+
+  The descriptor specifies the properties of the matrix in shared memory that 
+  is a multiplicand in the matrix multiply and accumulate operation. 
+  
+  The descriptor is a 64-bit value contained in a register with the following
+  ```
+  +---------+-----+-----------+-----+-----------+-----+-----+-----------+-----+
+  |   0-13  |14-15|   16-29   |30-31|   32-45   |46-48|49-51|   52-61   |62-63|
+  +---------+-----+-----------+-----+-----------+-----+-----+-----------+-----+
+  |  14bits |2bits|   14bits  |2bits|   14bits  |2bits|3bits|   10bits  |2bits|
+  +---------+-----+-----------+-----+-----------+-----+-----+-----------+-----+
+  | BaseAddr|  0  | LeadingDim|  0  |   Stride  |  0  |Offst|     0     | Mod |
+  +---------+-----+-----------+-----+-----------+-----+-----+-----------+-----+
+  ```
+   
+  See for more details:
+    https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#asynchronous-warpgroup-level-matrix-shared-memory-layout-matrix-descriptor
+
+  }];  
+  let results = (outs I64:$descriptor);
+  let arguments = (ins Arg<AnyMemRef, "", [MemWrite]>:$tensor, 
+                       NVGPU_TensorMapDescriptor:$tensorMap);
+  let assemblyFormat = [{$tensor `,` $tensorMap attr-dict `:` type($tensor) `,` type($tensorMap)}];
+  let hasVerifier = 1;
+}
+
 #endif // NVGPU

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

for (auto [index, value] : llvm::enumerate(coords)) {
       coords[index] = truncToI32(rewriter, op->getLoc(), value);
     }
 
     rewriter.replaceOpWithNewOp<NVVM::CpAsyncBulkTensorGlobalToSharedClusterOp>(
         op, dest, adaptor.getTensorMapDescriptor(), barrier, coords);
     return success();
   }
 };
+struct NVGPUGenerateGmmaDescriptorLowering
+    : public ConvertOpToLLVMPattern<nvgpu::GenerateGmmaDescriptorOp> {
+  using ConvertOpToLLVMPattern<
+      nvgpu::GenerateGmmaDescriptorOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(nvgpu::GenerateGmmaDescriptorOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op->getLoc();
+
+    nvgpu::TensorMapSwizzleKind shmemLayout =
+        op.getTensorMap().getType().getSwizzle();
+
+    MemRefType memrefType = op.getTensor().getType();
+    int64_t sizeH = memrefType.getDimSize(0);
+    // Note: currently only 128b swizzling is supported. Also not sure what to
+    // put without swizzling
+    unsigned swizzleSize =
+        (shmemLayout == nvgpu::TensorMapSwizzleKind::SWIZZLE_128B)  ? 128
+        : (shmemLayout == nvgpu::TensorMapSwizzleKind::SWIZZLE_64B) ? 64
+        : (shmemLayout == nvgpu::TensorMapSwizzleKind::SWIZZLE_32B) ? 32
+                                                                    : 1;
+
+    SmallVector<Value> extractedValue = getTypeConverter()->promoteOperands(
+        op->getLoc(), {op.getTensor()}, {adaptor.getTensor()}, rewriter);
+    if (extractedValue.empty())
+      return failure();
+
+    auto ti64 = rewriter.getIntegerType(64);
+    auto makeConst = [&](int32_t index) -> Value {
+      return rewriter.create<LLVM::ConstantOp>(
+          loc, ti64, rewriter.getI64IntegerAttr(index));
+    };
+    auto shiftLeft = [&](Value value, unsigned shift) -> Value {
+      return rewriter.create<LLVM::ShlOp>(loc, ti64, value, makeConst(shift));
+    };
+    auto shiftRight = [&](Value value, unsigned shift) -> Value {
+      return rewriter.create<LLVM::LShrOp>(loc, ti64, value, makeConst(shift));
+    };
+    auto makeOr = [&](Value lhs, Value rhs) -> Value {
+      return rewriter.create<LLVM::OrOp>(loc, ti64, lhs, rhs);
+    };
+
+    // Trim bits and get only 14 bits
+    Value basePtr = extractedValue.front();
+    Value ptri64 = rewriter.create<LLVM::PtrToIntOp>(loc, ti64, basePtr);
+    Value startAdress = shiftRight(shiftLeft(ptri64, 46), 50);
+    // Note: not sure this calculation is always correct.
+    Value strideDim = shiftRight(shiftLeft(makeConst(swizzleSize), 3), 4);
+    Value leadingDim = shiftRight(makeConst(swizzleSize * sizeH), 4);
+
+    Value descZero = makeConst(0);
+    // [62,64)  layout type
+    Value desc1 = makeOr(descZero, shiftLeft(makeConst(int(shmemLayout)), 62));
+    // [32,46)  stride
+    Value desc2 = makeOr(desc1, shiftLeft(strideDim, 32));
+    // [16,30)  leading dimension
+    Value desc3 = makeOr(desc2, shiftLeft(leadingDim, 16));
+    // [49,52)  base_offset
+    Value desc4 = makeOr(desc3, shiftLeft(makeConst(0), 49));
+    // [0,14)   start_address
+    Value desc = makeOr(desc4, startAdress);
+
+    rewriter.replaceOp(op, desc);
+    return success();
+  }
+};
 
 static Value makeI64Const(RewriterBase &rewriter, Operation *op,
                           int32_t index) {
   return rewriter.create<LLVM::ConstantOp>(op->getLoc(),
                                            rewriter.getIntegerType(64),
                                            rewriter.getI32IntegerAttr(index));
 }
 
patterns.add<
       NVGPUMBarrierArriveLowering,           // nvgpu.mbarrier.arrive
       NVGPUMBarrierArriveNoCompleteLowering, // nvgpu.mbarrier.arrive.no_complete
       NVGPUMBarrierTestWaitLowering,         // nvgpu.mbarrier.test_wait_parity
       NVGPUMBarrierTryWaitParityLowering,    // nvgpu.mbarrier.try_wait_parity
       NVGPUTmaAsyncLoadOpLowering,           // nvgpu.tma.async.load
       NVGPUTmaCreateDescriptorOpLowering,    // nvgpu.tma.create.descriptor
       NVGPUMBarrierArriveExpectTxLowering,   // nvgpu.mbarrier.arrive.expect_tx
       NVGPUTmaAsyncLoadOpLowering,           // nvgpu.tma.async.load
+      NVGPUGenerateGmmaDescriptorLowering,   // nvgpu.wgmma.generate.descriptor
       MmaSyncOptoNVVM, MmaLdMatrixOpToNVVM, NVGPUAsyncCopyLowering,
       NVGPUAsyncCreateGroupLowering, NVGPUAsyncWaitLowering,
       NVGPUMmaSparseSyncLowering>(converter);
 }

mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp

LogicalResult TmaCreateDescriptorOp::verify() {
   nvgpu::TensorMapDescriptorType desc = getTensorMap().getType();
   if (desc.getInterleave() != TensorMapInterleaveKind::INTERLEAVE_NONE)
     return emitError() << "Interleave options are not supported yet.";
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// NVGPU_GenerateGmmaDescriptorOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult GenerateGmmaDescriptorOp::verify() {
+  MemRefType memrefType = getTensor().getType();
+  MemRefType tensorMapType = getTensorMap().getType().getTensor();
+
+  if (memrefType != tensorMapType)
+    return emitError() << "memref and tensor map type mismatch";
+
+  if (!memrefType.hasStaticShape() || !tensorMapType.hasStaticShape())
+    return emitError() << "supports only static shapes";
+
+  if (memrefType.getRank() != 2)
+    return emitError() << "supports only 2d memref is supported for now";
+
+  if (getTensorMap().getType().getSwizzle() !=
+      TensorMapSwizzleKind::SWIZZLE_128B) {
+    return emitError() << "supports only "
+                       << stringifyTensorMapSwizzleKind(
+                              TensorMapSwizzleKind::SWIZZLE_128B)
+                       << " is supported for the time being";
+  }
+
+  if (getTensorMap().getType().getInterleave() !=
+      TensorMapInterleaveKind::INTERLEAVE_NONE) {
+    return emitError() << "supports only "
+                       << stringifyTensorMapInterleaveKind(
+                              TensorMapInterleaveKind::INTERLEAVE_NONE)
+                       << " is supported for the time being";
+  }
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd dialect, type, and op definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/NVGPU/IR/NVGPUAttrDefs.cpp.inc"
 
 #include "mlir/Dialect/NVGPU/IR/NVGPUEnums.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/NVGPU/IR/NVGPU.cpp.inc"
 
 #define GET_TYPEDEF_CLASSES
 #include "mlir/Dialect/NVGPU/IR/NVGPUTypes.cpp.inc"

mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir

-// RUN: mlir-opt --convert-nvgpu-to-nvvm='use-opaque-pointers=1' --split-input-file %s | FileCheck %s
+// RUN: mlir-opt --convert-nvgpu-to-nvvm='use-opaque-pointers=1' --split-input-file -cse -canonicalize %s | FileCheck %s
 
 // CHECK-LABEL: @m16n8k16_fp16
 func.func @m16n8k16_fp16(%arg0: vector<4x2xf16>, %arg1: vector<2x2xf16>, %arg2: vector<2x2xf16>) -> vector<2x2xf16> {
   // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<4 x vector<2xf16>>
   // CHECK: llvm.extractvalue %{{.*}}[1] : !llvm.array<4 x vector<2xf16>>
   // CHECK: llvm.extractvalue %{{.*}}[2] : !llvm.array<4 x vector<2xf16>>
   // CHECK: llvm.extractvalue %{{.*}}[3] : !llvm.array<4 x vector<2xf16>>
   // CHECK: llvm.extractvalue %{{.*}}[0] : !llvm.array<2 x vector<2xf16>>
func.func @create_tensor_map(%devicePtr2d : memref<64x128xf32>, %devicePtr1d : memref<128xf32>) {
   // CHECK : llvm.call @mgpuTensorMapEncodeTiledMemref
   %tensorMap2d = nvgpu.tma.create.descriptor %devicePtr2d_unranked box[%crd0, %crd1] : memref<*xf32> -> !tensorMap2d
 
   %devicePtr1d_unranked = memref.cast %devicePtr1d : memref<128xf32> to memref<*xf32>
   // CHECK : llvm.call @mgpuTensorMapEncodeTiledMemref
   %tensorMap1d = nvgpu.tma.create.descriptor %devicePtr1d_unranked box[%crd1] : memref<*xf32> -> !tensorMap1d
   func.return
 }
+
+// -----
+
+!tensorMap = !nvgpu.tensormap.descriptor<tensor = memref<128x64xf16,3>, swizzle = swizzle_128b, l2promo=none, oob=zero, interleave=none>
+memref.global "private" @dynamicShmem : memref<0xf16,3>
+func.func @create_wgmma_descriptor(%tensorMap : !tensorMap) -> i64{
+  %dynamicMem = memref.get_global @dynamicShmem : memref<0xf16, 3>
+  %lhsShmem = memref.reinterpret_cast %dynamicMem to offset: [0], sizes: [128,64], strides: [64,1] : memref<0xf16, 3> to memref<128x64xf16,3>
+  // CHECK: %[[S0:.+]] = memref.get_global @dynamicShmem : memref<0xf16, 3>
+  // CHECK: %[[R0:.+]] = memref.reinterpret_cast %[[S0]] to offset: [0], sizes: [128, 64], strides: [64, 1] : memref<0xf16, 3> to memref<128x64xf16, 3>
+  // CHECK: %[[S1:.+]] = builtin.unrealized_conversion_cast %[[R0]] : memref<128x64xf16, 3> to !llvm.struct<(ptr<3>, ptr<3>, i64, array<2 x i64>, array<2 x i64>)>
+  // CHECK: %[[S2:.+]] = llvm.extractvalue %[[S1]][0] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<2 x i64>, array<2 x i64>)>
+  // CHECK: %[[PTR:.+]] = llvm.ptrtoint %[[S2]] : !llvm.ptr<3> to i64
+  // CHECK: %[[S10:.+]] = llvm.mlir.constant(46 : i64) : i64
+  // CHECK: %[[S11:.+]] = llvm.shl %[[PTR]], %[[S10]]  : i64
+  // CHECK: %[[S12:.+]] = llvm.mlir.constant(50 : i64) : i64
+  // CHECK: %[[S13:.+]] = llvm.lshr %[[S11]], %[[S12]]  : i64
+  // CHECK: %[[S14:.+]] = llvm.mlir.constant(128 : i64) : i64
+  // CHECK: %[[S15:.+]] = llvm.mlir.constant(3 : i64) : i64
+  // CHECK: %[[S16:.+]] = llvm.shl %[[S14]], %[[S15]]  : i64
+  // CHECK: %[[S17:.+]] = llvm.mlir.constant(4 : i64) : i64
+  // CHECK: %[[S18:.+]] = llvm.lshr %[[S16]], %[[S17]]  : i64
+  // CHECK: %[[S19:.+]] = llvm.mlir.constant(16384 : i64) : i64
+  // CHECK: %[[S20:.+]] = llvm.mlir.constant(4 : i64) : i64
+  // CHECK: %[[S21:.+]] = llvm.lshr %[[S19]], %[[S20]]  : i64
+  // CHECK: %[[S22:.+]] = llvm.mlir.constant(0 : i64) : i64
+  // CHECK: %[[S23:.+]] = llvm.mlir.constant(3 : i64) : i64
+  // CHECK: %[[S24:.+]] = llvm.mlir.constant(62 : i64) : i64
+  // CHECK: %[[S25:.+]] = llvm.shl %[[S23]], %[[S24]]  : i64
+  // CHECK: %[[S26:.+]] = llvm.or %[[S22]], %[[S25]]  : i64
+  // CHECK: %[[S27:.+]] = llvm.mlir.constant(32 : i64) : i64
+  // CHECK: %[[S28:.+]] = llvm.shl %[[S18]], %[[S27]]  : i64
+  // CHECK: %[[S29:.+]] = llvm.or %[[S26]], %[[S28]]  : i64
+  // CHECK: %[[S30:.+]] = llvm.mlir.constant(16 : i64) : i64
+  // CHECK: %[[S31:.+]] = llvm.shl %[[S21]], %[[S30]]  : i64
+  // CHECK: %[[S32:.+]] = llvm.or %[[S29]], %[[S31]]  : i64
+  // CHECK: %[[S33:.+]] = llvm.mlir.constant(0 : i64) : i64
+  // CHECK: %[[S34:.+]] = llvm.mlir.constant(49 : i64) : i64
+  // CHECK: %[[S35:.+]] = llvm.shl %[[S33]], %[[S34]]  : i64
+  // CHECK: %[[S36:.+]] = llvm.or %[[S32]], %[[S35]]  : i64
+  // CHECK: %[[DESC:.+]] = llvm.or %[[S36]], %[[S13]]  : i64
+  // CHECK: return %[[DESC]] : i64
+  %descA = nvgpu.wgmma.generate.descriptor %lhsShmem, %tensorMap : memref<128x64xf16,3>, !tensorMap
+  func.return %descA : i64
+}