[mlir][nvgpu] NFC - move NVGPU conversion helpers to NvGpu utils library

[christopherbate]

Summary

The ConvertVectorToGpu pass implementation contained a small private
support library for performing various calculations during conversion
between vector and nvgpu.mma.sync and nvgpu.ldmatrix operations.
The support library is moved under Dialect/NVGPU/Utils because the
functions have wider utility. Some documentation comments are added or
improved.

[aartbik]

mlir/lib/Dialect/NVGPU/Utils/CMakeLists.txt

1

This will need a bazel update.

mlir/include/mlir/Dialect/NVGPU/Utils/MMAUtils.h

-//===- NvvmMMASupport.h - MLIR Vector to GPU lowering support --------===//
+//===-- MMAUtils.h - MLIR NVGPU dialect utilities for MMA operations-------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
-// This file provides utilities to assist in the lowering of Vector operations
-// to GPU dialect MMA operations.
+// This file provides utilities to assist in the lowering of other dialects
+// (e.g. Vector) to `nvgpu.mma.*` dialect operations.
 //
 //===----------------------------------------------------------------------===//
-#ifndef MLIR_CONVERSION_VECTORTOGPU_NVGPUSUPPORT_H
-#define MLIR_CONVERSION_VECTORTOGPU_NVGPUSUPPORT_H
+#ifndef MLIR_DIALECT_NVGPU_UTILS_MMAUTILS_H
+#define MLIR_DIALECT_NVGPU_UTILS_MMAUTILS_H
 
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
-#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
-#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
-#include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Types.h"
 
 namespace mlir {
+namespace vector {
+enum class IteratorType : uint32_t;
+class ContractionOp;
+} // namespace vector
+
+namespace NVVM {
+enum class MMALayout : uint32_t;
+} // namespace NVVM
+
 namespace nvgpu {
 
+/// Represents the role of an operand in an MMA instruction:
+/// `result := matmul(A, B) + C`
 enum class MatMulOperandRole : int32_t { A = 0, B, C };
 
 /// Collects information about a warp-level matrix operand represented by a
 /// VectorType.
 struct WarpMatrixInfo {
   VectorType vectorType;
   MatMulOperandRole operandRole;
 };
 
-/// Given an op that operates on a VectorType representing a warp-level matrix
-/// operand, the function returns a struct containing relevant type information.
+/// If `op` is a `vector.transfer_write`, return the `WarpMatrixInfo` for the
+/// vector operand. If op is a `vector.transfer_read`, `vector.contraction`, or
+/// `arith.constant`, return the `WarpMatrixInfo` corresponding to the result.
+/// Otherwise, return failure.
 FailureOr<WarpMatrixInfo> getWarpMatrixInfo(Operation *op);
 
 /// Returns the number of bits in a single tile row. It is either 128, 256, or
 /// 512 bits depending on the data type and` whether the operand is an
 /// accumulator/result operand
 int64_t inferTileWidthInBits(const WarpMatrixInfo &type);
 
 /// Specifies information about the registers which compose a matrix fragment

 /// for (AKA the matrix fragment values) during a warp-collective matrix
 /// operation. For a visual reference of this LaneId -> (row, col) mapping,
 /// please see NVIDIA's PTX documentation:
 /// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-for-mma
 FailureOr<AffineMap>
 getLaneIdAndValueIdToOperandCoord(Location loc, OpBuilder &builder,
                                   const WarpMatrixInfo &fragmentType);
 
+/// Encapsulates the parameters needed to lower a `nvgpu.ldmatrix` operation to
+/// `nvvm.ldmatrix`.
 struct LdMatrixParams {
   VectorType fragmentType;
   bool isAccum;
   int64_t numTiles;
   vector::IteratorType contiguousDimType;
   NVVM::MMALayout targetLayout;
 };
 
+/// Given `type` that contains info for a warp-matrix operand and whether or not
+/// the load is a transposed load, return the LdMatrixParams.
 FailureOr<LdMatrixParams> getLdMatrixParams(const WarpMatrixInfo &type,
                                             bool transpose);
 /// Returns an AffineMap which maps a single dimension representing the laneId
 /// to two results representing offsets within the matrix operand that should
 /// be the pointer locations a thread should pass to the ldmatrix instruction.
 FailureOr<AffineMap>
 getLaneIdToLdMatrixMatrixCoord(Location loc, OpBuilder &builder,
                                const LdMatrixParams &params);
 
-// Transform contract into (m, k)x(n, k)x(m, n) form so that it can be converted
-// to MMA matmul.
+/// Transform `vector.contract` into (m,k)x(n,k)x(m,n) form so that it can be
+/// converted to `nvgpu.mma.sync`. This specific form is meant to indicate that
+/// the vector operands are organized such that the reduction dimension is
+/// contiguous.
 struct PrepareContractToGPUMMASync
     : public OpRewritePattern<vector::ContractionOp> {
   using OpRewritePattern<vector::ContractionOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::ContractionOp op,
                                 PatternRewriter &rewriter) const override;
 };
 
 } // namespace nvgpu
 } // namespace mlir
 
-#endif // MLIR_CONVERSION_VECTORTOGPU_NVGPUSUPPORT_H
+#endif // MLIR_DIALECT_NVGPU_UTILS_MMAUTILS_H

mlir/lib/Conversion/VectorToGPU/CMakeLists.txt

 add_mlir_conversion_library(MLIRVectorToGPU
   VectorToGPU.cpp
-  NvGpuSupport.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/VectorToGPU
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRGPUOps
   MLIRLLVMDialect
   MLIRMemRefDialect
   MLIRNVGPUDialect
+  MLIRNVGPUUtils
   MLIRTransforms
   MLIRVectorDialect
   MLIRVectorUtils
   )

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

 //===- VectorToGPU.cpp - Convert vector to GPU dialect ----------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements lowering of vector operations to GPU dialect ops.
 //
 //===----------------------------------------------------------------------===//
 
-#include <type_traits>
-
-#include "NvGpuSupport.h"
 #include "mlir/Conversion/VectorToGPU/VectorToGPU.h"
 
+#include <type_traits>
+
 #include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
+#include "mlir/Dialect/NVGPU/Utils/MMAUtils.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"

mlir/lib/Dialect/NVGPU/CMakeLists.txt

 add_subdirectory(IR)
+add_subdirectory(Utils)
 add_subdirectory(Transforms)

mlir/lib/Dialect/NVGPU/Utils/CMakeLists.txt

+add_mlir_dialect_library(MLIRNVGPUUtils
+  MMAUtils.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arithmetic/Utils
+
+  LINK_LIBS PUBLIC
+  MLIRAffineDialect
+  MLIRLLVMDialect
+  MLIRNVGPUDialect
+  MLIRNVVMDialect
+  MLIRVectorDialect
+  MLIRIR
+  )

mlir/lib/Dialect/NVGPU/Utils/MMAUtils.cpp

-//===- NvGpuSupport.cpp - MLIR Vector to GPU lowering support --------===//
+//===- MMAUtils.cpp - MLIR NVGPU dialect utils for MMA operations----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
-//
-// This file provides utilities to assist in the lowering of Vector operations
-// to NvGPU dialect MMA operations.
-//
-//===----------------------------------------------------------------------===//
+#include "mlir/Dialect/NVGPU/Utils/MMAUtils.h"
 
-#include "NvGpuSupport.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 
-namespace mlir {
-namespace nvgpu {
-namespace {
+using namespace mlir;
+using namespace mlir::nvgpu;
 
 /// There are always 4 threads per [128|256|512] bit row.
-constexpr int64_t kThreadsPerRow = 4;
-
-constexpr int64_t kNumRowsPerTile = 8;
+static constexpr int64_t kThreadsPerRow = 4;
+static constexpr int64_t kNumRowsPerTile = 8;
 
-bool isAccumulatorOrResult(MatMulOperandRole operandType) {
+static bool isAccumulatorOrResult(MatMulOperandRole operandType) {
   return operandType == MatMulOperandRole::C;
 }
 
 /// Returns the number of registers which compose a matrix fragment held by a
 /// single thread.
-int64_t inferNumRegistersPerMatrixFragment(const WarpMatrixInfo &type) {
+static int64_t inferNumRegistersPerMatrixFragment(const WarpMatrixInfo &type) {
   int64_t lineSize = inferTileWidthInBits(type);
   auto shape = type.vectorType.getShape();
   return (shape[0] / kNumRowsPerTile) *
          (shape[1] * type.vectorType.getElementType().getIntOrFloatBitWidth()) /
          lineSize;
 }
 
 /// Returns the number of 8 x [128|256|512] bit tiles that compose the given
 /// operand shape.
-std::array<int64_t, 2> getTileShape(ArrayRef<int64_t> operandShape,
-                                    Type elementType, int64_t lineSizeBits) {
+static std::array<int64_t, 2> getTileShape(ArrayRef<int64_t> operandShape,
+                                           Type elementType,
+                                           int64_t lineSizeBits) {
   // For each 8x128bit square, a thread is responsible for one 32bit register.
   return {operandShape[0] / kNumRowsPerTile,
           (operandShape[1] * elementType.getIntOrFloatBitWidth()) /
               lineSizeBits};
 }
 
-} // namespace
-
-FailureOr<WarpMatrixInfo> getWarpMatrixInfo(Operation *op) {
+FailureOr<WarpMatrixInfo> nvgpu::getWarpMatrixInfo(Operation *op) {
   WarpMatrixInfo info;
 
   // Determine the vector type.
   if (vector::TransferWriteOp writeOp = dyn_cast<vector::TransferWriteOp>(op)) {
     info.vectorType = writeOp.getVectorType();
   } else if (isa<vector::TransferReadOp, vector::ContractionOp,
                  arith::ConstantOp>(op)) {
     info.vectorType = op->getResult(0).getType().cast<VectorType>();
for (Operation *user : op->getUsers()) {
     if (contract.getRhs() == op->getResult(0)) {
       info.operandRole = MatMulOperandRole::B;
       break;
     }
   }
   return info;
 }
 
-int64_t inferTileWidthInBits(const WarpMatrixInfo &type) {
+int64_t nvgpu::inferTileWidthInBits(const WarpMatrixInfo &type) {
   bool isAcc = isAccumulatorOrResult(type.operandRole);
   Type elType = type.vectorType.getElementType();
   if (isAcc && elType.getIntOrFloatBitWidth() == 32) {
     return 256;
   }
   if (elType.getIntOrFloatBitWidth() == 64) {
     return isAcc ? 512 : 256;
   }
   return 128;
 }
 
 FailureOr<FragmentElementInfo>
-getMmaSyncRegisterType(const WarpMatrixInfo &type) {
+nvgpu::getMmaSyncRegisterType(const WarpMatrixInfo &type) {
   MLIRContext *ctx = type.vectorType.getContext();
   const bool isAccum = isAccumulatorOrResult(type.operandRole);
 
   Type elType = type.vectorType.getElementType();
   if (elType.isF16()) {
     return FragmentElementInfo{
         LLVM::getFixedVectorType(Float16Type::get(ctx), 2), 2, 32,
         inferNumRegistersPerMatrixFragment(type)};
static AffineMap getRegisterIndexToTileOffsetMap(int64_t lineSize,
   return AffineMap::get(
       2, 0,
       {(registerIdx % num8x128bTiles[0]) * 8,
        (registerIdx.floorDiv(num8x128bTiles[0])) * elementsPerLine},
       elementType.getContext());
 }
 
 FailureOr<AffineMap>
-getLaneIdAndValueIdToOperandCoord(Location loc, OpBuilder &builder,
+nvgpu::getLaneIdAndValueIdToOperandCoord(Location loc, OpBuilder &builder,
                                          const WarpMatrixInfo &fragmentType) {
   Type elementType = fragmentType.vectorType.getElementType();
   ArrayRef<int64_t> operandShape = fragmentType.vectorType.getShape();
   FailureOr<nvgpu::FragmentElementInfo> regInfo =
       getMmaSyncRegisterType(fragmentType);
   if (failed(regInfo))
     return failure();
 
   const int64_t elementBitWidth = elementType.getIntOrFloatBitWidth();
nvgpu::getLaneIdAndValueIdToOperandCoord(Location loc, OpBuilder &builder,
 
   auto tileRow = registerIndexToTileCoord.getResult(0);
   auto tileCol = registerIndexToTileCoord.getResult(1);
   return makeMap({tileRow + laneId.floorDiv(kThreadsPerRow),
                   tileCol + (laneId % kThreadsPerRow) * elementsPerRegister +
                       (logicalValueIdDim % elementsPerRegister)});
 }
 
-FailureOr<nvgpu::LdMatrixParams> getLdMatrixParams(const WarpMatrixInfo &type,
-                                                   bool transpose) {
+FailureOr<nvgpu::LdMatrixParams>
+nvgpu::getLdMatrixParams(const WarpMatrixInfo &type, bool transpose) {
   LdMatrixParams params;
   Type elType = type.vectorType.getElementType();
   params.fragmentType = type.vectorType;
   if (type.operandRole == MatMulOperandRole::A ||
       type.operandRole == MatMulOperandRole::C) {
     params.targetLayout = NVVM::MMALayout::row;
   } else {
     params.targetLayout = NVVM::MMALayout::col;
nvgpu::getLdMatrixParams(const WarpMatrixInfo &type, bool transpose) {
 
   if (params.numTiles == 0)
     return failure();
 
   return params;
 }
 
 FailureOr<AffineMap>
-getLaneIdToLdMatrixMatrixCoord(Location loc, OpBuilder &builder,
+nvgpu::getLaneIdToLdMatrixMatrixCoord(Location loc, OpBuilder &builder,
                                       const LdMatrixParams &params) {
   // One thread per 128b row.
   const int64_t kNumThreadsPerTile = kNumRowsPerTile;
   const int bitsPerElement = static_cast<int>(
       params.fragmentType.getElementType().getIntOrFloatBitWidth());
   const int kElementsPer128b = (128 / bitsPerElement);
   ArrayRef<int64_t> operandShape = params.fragmentType.getShape();
   AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
 
if (params.contiguousDimType == vector::IteratorType::parallel) {
     auto tileCol = (groupIdx % num8x128bCols);
     auto tileRow = groupIdx.floorDiv(num8x128bCols);
     return makeMap({tileCol * kElementsPer128b,
                     tileRow * kNumRowsPerTile + (d0 % kNumRowsPerTile)});
   }
   return failure();
 }
 
-LogicalResult
-PrepareContractToGPUMMASync::matchAndRewrite(vector::ContractionOp op,
-                                             PatternRewriter &rewriter) const {
+LogicalResult nvgpu::PrepareContractToGPUMMASync::matchAndRewrite(
+    vector::ContractionOp op, PatternRewriter &rewriter) const {
   Location loc = op.getLoc();
   Value lhs = op.getLhs();
   Value rhs = op.getRhs();
   Value res = op.getAcc();
 
   // Set up the parallel/reduction structure in right form.
   using MapList = ArrayRef<ArrayRef<AffineExpr>>;
   auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
LogicalResult nvgpu::PrepareContractToGPUMMASync::matchAndRewrite(
   } else {
     return failure();
   }
   rewriter.replaceOpWithNewOp<vector::ContractionOp>(
       op, lhs, rhs, res, rewriter.getAffineMapArrayAttr(canonicalForm),
       op.getIteratorTypes());
   return success();
 }
-
-} // namespace nvgpu
-} // namespace mlir