Index: mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp =================================================================== --- mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp +++ mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp @@ -650,6 +650,34 @@ return success(); } +/// Check if the loaded matrix operand requires transposed. +/// Transposed Map Example: +/// Example 1 : (..., d0, d1) -> (d1 * 1, d0 * 2) +/// Example 2 : (d0, d1, d2, d3) -> (d3, d2) +/// +/// The code below checks if the output 2D is transposed using a generalized +/// version : (d0, d1, dn, ..., dm, ...) -> (dm, dn) +/// Returns : true; if m > n, false o.w. + +static bool isTransposed(vector::TransferReadOp op) { + mlir::AffineMap map = op.getPermutationMap(); + if (map.getNumResults() != 2) { + op->emitError("Expected 2D transfer read"); + } + + // Output 2D matrix dimensions in the order of d0, d1. + auto dM = map.getResult(0); + auto dN = map.getResult(1); + + // Find the position of these expressions in the input. + auto exprM = dM.dyn_cast(); + auto exprN = dN.dyn_cast(); + if (!exprM || !exprN) { + op->emitError("Expected to find AffineDimExpr in vector::TransferReadOp"); + } + return exprM.getPosition() > exprN.getPosition(); +} + static LogicalResult creatLdMatrixCompatibleLoads(RewriterBase &rewriter, vector::TransferReadOp op, llvm::DenseMap &valueMapping) { @@ -671,9 +699,10 @@ return rewriter.notifyMatchFailure(op, "not mma sync reg info"); } - FailureOr params = nvgpu::getLdMatrixParams( - *warpMatrixInfo, - /*transpose=*/!op.getPermutationMap().isMinorIdentity()); + FailureOr params = + nvgpu::getLdMatrixParams(*warpMatrixInfo, + /*transpose=*/isTransposed(op)); + if (failed(params)) { LLVM_DEBUG( DBGS() @@ -700,7 +729,7 @@ indices); nvgpu::LdMatrixOp newOp = rewriter.create( loc, vectorType, op.getSource(), indices, - !op.getPermutationMap().isMinorIdentity(), params->numTiles); + /*transpose=*/isTransposed(op), params->numTiles); valueMapping[op] = newOp->getResult(0); return success(); } Index: mlir/test/Conversion/VectorToGPU/vector-to-mma-ops-mma-sync.mlir =================================================================== --- mlir/test/Conversion/VectorToGPU/vector-to-mma-ops-mma-sync.mlir +++ mlir/test/Conversion/VectorToGPU/vector-to-mma-ops-mma-sync.mlir @@ -253,6 +253,63 @@ } // ----- +//##################################################################################################### +// FP16 row-row-row (split-k with multi-buffering, ldmatrix x4 for matrixA and ldmatrix x4 for matrixB) +//##################################################################################################### + +// CHECK-DAG: [[$strided_map:#.+]] = affine_map<()[s0] -> (s0 mod 16)> +// CHECK-DAG: [[$contiguous_map:#.+]] = affine_map<()[s0] -> ((s0 floordiv 16) * 8)> + +#map0 = affine_map<(d0, d1, d2) -> (d2, d1)> +#map1 = affine_map<(d0, d1, d2) -> (d0, d2)> +#map2 = affine_map<(d0, d1, d2) -> (d1, d2)> +#map3 = affine_map<(d0, d1, d2) -> (d0, d1)> +#map_a = affine_map<(d0, d1, d2, d3) -> (d1, d3)> +#map_b = affine_map<(d0, d1, d2, d3) -> (d3, d2)> + + +// CHECK-LABEL: func @splitk_multibuffering_m16n8k16_fp16_row_row_row +func.func @splitk_multibuffering_m16n8k16_fp16_row_row_row(%arg0: memref<4x32x1x32xf16, #gpu.address_space>, %arg1: memref<4x1x32x32xf16, #gpu.address_space>, %arg2: memref<1x32x40xf16, #gpu.address_space>) { + + %c0 = arith.constant 0 : index + %cst = arith.constant 0.000000e+00 : f16 + + // CHECK-DAG: [[m_coord:%.+]] = affine.apply [[$strided_map]] + // CHECK-DAG: [[k_coord:%.+]] = affine.apply [[$contiguous_map]] + // CHECK: [[fragmentA:%.+]] = nvgpu.ldmatrix %arg0[%c0, [[m_coord]], %c0, [[k_coord]]] {numTiles = 4 : i32, transpose = false} + %A = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %cst {in_bounds = [true, true], permutation_map = #map_a} : memref<4x32x1x32xf16, #gpu.address_space>, vector<16x16xf16> + + // CHECK-DAG: [[n_coord:%.+]] = affine.apply [[$contiguous_map]] + // CHECK-DAG: [[k_coord:%.+]] = affine.apply [[$strided_map]] + // CHECK-DAG: [[fragmentB:%.+]] = nvgpu.ldmatrix %arg1[%c0, %c0, [[k_coord]], [[n_coord]]] {numTiles = 4 : i32, transpose = true} + %B = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %cst {in_bounds = [true, true], permutation_map = #map_b} : memref<4x1x32x32xf16, #gpu.address_space>, vector<16x16xf16> + + // CHECK-DAG: [[m_coord:%.+]] = affine.apply [[$strided_map]] + // CHECK-DAG: [[n_coord:%.+]] = affine.apply [[$contiguous_map]] + // CHECK-DAG: [[fragmentC:%.*]] = nvgpu.ldmatrix %arg2[%c0, [[m_coord]], [[n_coord]]] {numTiles = 4 : i32, transpose = false} + %C = vector.transfer_read %arg2[%c0, %c0, %c0], %cst {in_bounds = [true, true]} : memref<1x32x40xf16, #gpu.address_space>, vector<16x16xf16> + + // CHECK-DAG: [[fragmentB0:%.+]] = vector.extract_strided_slice [[fragmentB]] {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x2xf16> to vector<2x2xf16> + // CHECK-DAG: [[fragmentC0:%.+]] = vector.extract_strided_slice [[fragmentC]] {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x2xf16> to vector<2x2xf16> + // CHECK: nvgpu.mma.sync([[fragmentA]], [[fragmentB0]], [[fragmentC0]]) {mmaShape = [16, 8, 16]} : (vector<4x2xf16>, vector<2x2xf16>, vector<2x2xf16>) -> vector<2x2xf16> + %B0 = vector.extract_strided_slice %B {offsets = [0, 0], sizes = [8, 16], strides = [1, 1]} : vector<16x16xf16> to vector<8x16xf16> + %C0 = vector.extract_strided_slice %C {offsets = [0, 0], sizes = [16, 8], strides = [1, 1]} : vector<16x16xf16> to vector<16x8xf16> + %D0 = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind} %A, %B0, %C0 : vector<16x16xf16>, vector<8x16xf16> into vector<16x8xf16> + vector.transfer_write %D0, %arg2[%c0, %c0, %c0] {in_bounds = [true, true]} : vector<16x8xf16>, memref<1x32x40xf16, #gpu.address_space> + + // CHECK-DAG: [[fragmentB1:%.+]] = vector.extract_strided_slice [[fragmentB]] {offsets = [2, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x2xf16> to vector<2x2xf16> + // CHECK-DAG: [[fragmentC1:%.+]] = vector.extract_strided_slice [[fragmentC]] {offsets = [2, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x2xf16> to vector<2x2xf16> + // CHECK: nvgpu.mma.sync([[fragmentA]], [[fragmentB1]], [[fragmentC1]]) {mmaShape = [16, 8, 16]} : (vector<4x2xf16>, vector<2x2xf16>, vector<2x2xf16>) -> vector<2x2xf16> + %B1 = vector.extract_strided_slice %B {offsets = [8, 0], sizes = [8, 16], strides = [1, 1]} : vector<16x16xf16> to vector<8x16xf16> + %C1 = vector.extract_strided_slice %C {offsets = [0, 8], sizes = [16, 8], strides = [1, 1]} : vector<16x16xf16> to vector<16x8xf16> + %D1 = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind} %A, %B1, %C1 : vector<16x16xf16>, vector<8x16xf16> into vector<16x8xf16> + vector.transfer_write %D1, %arg2[%c0, %c0, %c0] {in_bounds = [true, true]} : vector<16x8xf16>, memref<1x32x40xf16, #gpu.address_space> + + return +} + +// ----- + // CHECK-DAG: [[$strided_map:#.+]] = affine_map<()[s0] -> (s0 mod 16)> // CHECK-DAG: [[$contiguous_map:#.+]] = affine_map<()[s0] -> ((s0 floordiv 16) * 8)>