[mlir][VectorToGPU] Fix support for i4, col-major operand support

[christopherbate]

Summary

For the conversion to nvgpu mma.sync and ldmatrix pathways, the code
was missing support for the i4 data type. While fixing this, another
bug was discoverd that caused the number of ldmatrix tiles calculated for
certain operand types and configurations to be incorrect. This change
fixes both issues and adds additional tests.

[christopherbate]

Add fix/test for tf32 col/col/row

[christopherbate]

Rebase

mlir/lib/Conversion/VectorToGPU/NvGpuSupport.cpp

getMmaSyncRegisterType(const WarpMatrixInfo &type) {
   }
 
   // int8 operand
   if (elType.isInteger(8)) {
     return FragmentElementInfo{
         LLVM::getFixedVectorType(IntegerType::get(ctx, 8), 4), 4, 32,
         inferNumRegistersPerMatrixFragment(type)};
   }
+
+  // int4 operand
+  if (elType.isInteger(4)) {
+    return FragmentElementInfo{
+        LLVM::getFixedVectorType(IntegerType::get(ctx, 4), 8), 8, 32,
+        inferNumRegistersPerMatrixFragment(type)};
+  }
+
   // Integer 32bit acc operands
   if (elType.isInteger(32)) {
     return FragmentElementInfo{
         LLVM::getFixedVectorType(IntegerType::get(ctx, 32), 2), 2, 64,
         inferNumRegistersPerMatrixFragment(type)};
   }
 
   // Floating point 32bit operands
if (type.operandRole == MatMulOperandRole::A ||
     params.targetLayout = NVVM::MMALayout::row;
   } else {
     params.targetLayout = NVVM::MMALayout::col;
   }
   ArrayRef<int64_t> shape = type.vectorType.getShape();
   params.contiguousDimType =
       transpose ? IteratorType::Parallel : IteratorType::Reduction;
 
-  if (params.targetLayout == NVVM::MMALayout::row) {
+  if (params.contiguousDimType == IteratorType::Reduction) {
     params.numTiles = (shape[0] / kNumRowsPerTile) *
                       ((shape[1] * elType.getIntOrFloatBitWidth()) / 128);
   } else {
     params.numTiles = (shape[1] / kNumRowsPerTile) *
                       ((shape[0] * elType.getIntOrFloatBitWidth()) / 128);
   }
 
   if (params.numTiles == 0)
PrepareContractToGPUMMASync::matchAndRewrite(vector::ContractionOp op,
   rewriter.replaceOpWithNewOp<vector::ContractionOp>(
       op, lhs, rhs, res, rewriter.getAffineMapArrayAttr(canonicalForm),
       op.getIteratorTypes());
   return success();
 }
 
 } // namespace nvgpu
 } // namespace mlir
 No newline at end of file

mlir/test/Conversion/VectorToGPU/vector-to-mma-ops-mma-sync.mlir

func.func @m16n8k8_tf32_f32_row_row_row(%arg0: memref<20x20xf32, 3>, %arg1: memref<20x20xf32, 3>, %arg2: memref<20x20xf32>) {
   // CHECK: vector.store
   // CHECK: vector.extract [[d_frag]][1] : vector<2x2xf32>
   // CHECK: affine.apply [[$rowC8_map]]
   // CHECK: affine.apply [[$colC_map]]
   // CHECK: vector.store
   vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x8xf32>, memref<20x20xf32>
   return
 }
+
+// -----
+
+//#########################################################
+// INT4 row-col-row
+//#########################################################
+
+// CHECK-DAG: [[$rowA0_map:#.+]] = affine_map<()[s0] -> (s0 mod 16)>
+// CHECK-DAG: [[$colA0_map:#.+]] = affine_map<()[s0] -> ((s0 floordiv 16) * 32)>
+// CHECK-DAG: [[$rowB0_map:#.+]] = affine_map<()[s0] -> (s0 mod 8)>
+// CHECK-DAG: [[$colB0_map:#.+]] = affine_map<()[s0] -> ((s0 floordiv 8) * 32)>
+// CHECK-DAG: [[$rowC0_map:#.+]] = affine_map<()[s0] -> (s0 floordiv 4)>
+// CHECK-DAG: [[$colC0_map:#.+]] = affine_map<()[s0] -> (s0 * 2 - (s0 floordiv 4) * 8
+// CHECK-DAG: [[$rowC8_map:#.+]] = affine_map<()[s0] -> (s0 floordiv 4 + 8)>
+
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: func @m16n8k64_int4_row_col_row
+func.func @m16n8k64_int4_row_col_row(%arg0: memref<128x128xi4, 3>, %arg1: memref<128x128xi4, 3>, %arg2: memref<128x128xi32>) {
+  %cst  = arith.constant 0 : i4
+  %cst0  = arith.constant 0 : i32
+  %cst_0 = arith.constant dense<0> : vector<32x8xi4>
+  %c0 = arith.constant 0 : index
+
+  // CHECK: [[lane:%.+]] = gpu.lane_id
+  // CHECK: [[row:%.+]] = affine.apply [[$rowA0_map]]()[[[lane]]]
+  // CHECK: [[col:%.+]] = affine.apply [[$colA0_map]]()[[[lane]]]
+  // CHECK: nvgpu.ldmatrix %arg0[[[row]], [[col]]] {numTiles = 4 : i32, transpose = false} : memref<128x128xi4, 3> -> vector<4x8xi4>
+
+  // CHECK: [[lane:%.+]] = gpu.lane_id
+  // CHECK: [[row:%.+]] = affine.apply [[$rowB0_map]]()[[[lane]]]
+  // CHECK: [[col:%.+]] = affine.apply [[$colB0_map]]()[[[lane]]]
+  // CHECK: nvgpu.ldmatrix %arg1[[[row]], [[col]]] {numTiles = 2 : i32, transpose = false} : memref<128x128xi4, 3> -> vector<2x8xi4>
+  
+  // CHECK: [[lane:%.+]] = gpu.lane_id
+  // CHECK: [[row:%.+]] = affine.apply [[$rowC0_map]]()[{{%.+}}]
+  // CHECK: [[col:%.+]] = affine.apply [[$colC0_map]]()[{{%.+}}]
+  // CHECK: vector.load %arg2[[[row]], [[col]]] : memref<128x128xi32>, vector<2xi32>
+
+  // CHECK: [[row:%.+]] = affine.apply [[$rowC8_map]]()[{{%.+}}]
+  // CHECK: [[col:%.+]] = affine.apply [[$colC0_map]]()[{{%.+}}]
+  // CHECK: vector.load %arg2[[[row]], [[col]]] : memref<128x128xi32>, vector<2xi32>
+  // CHECK-NOT: vector.load
+
+  %A = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : memref<128x128xi4, 3>, vector<16x64xi4>
+  %B = vector.transfer_read %arg1[%c0, %c0], %cst {in_bounds = [true, true]} : memref<128x128xi4, 3>, vector<8x64xi4>
+  %C = vector.transfer_read %arg2[%c0, %c0], %cst0 {in_bounds = [true, true]} : memref<128x128xi32>, vector<16x8xi32>
+  // CHECK: [[d:%.+]] = nvgpu.mma.sync({{.*}}) {mmaShape = [16, 8, 64]} : (vector<4x8xi4>, vector<2x8xi4>, vector<2x2xi32>) -> vector<2x2xi32>
+  %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %A, %B, %C : vector<16x64xi4>, vector<8x64xi4> into vector<16x8xi32>
+
+  // CHECK: [[lane:%.+]] = gpu.lane_id
+  // CHECK: [[v:%.+]] = vector.extract [[d]][0] : vector<2x2xi32>  
+  // CHECK: [[row:%.+]] = affine.apply [[$rowC0_map]]()[[[lane]]]
+  // CHECK: [[col:%.+]] = affine.apply [[$colC0_map]]()[[[lane]]]
+  // CHECK: vector.store [[v]], %arg2[[[row]], [[col]]] : memref<128x128xi32>, vector<2xi32>
+  
+  // CHECK: [[v:%.+]] = vector.extract [[d]][1] : vector<2x2xi32>
+  // CHECK: [[row:%.+]] = affine.apply [[$rowC8_map]]()[[[lane]]]
+  // CHECK: [[col:%.+]] = affine.apply [[$colC0_map]]()[[[lane]]]
+  // CHECK: vector.store [[v]], %arg2[[[row]], [[col]]] : memref<128x128xi32>, vector<2xi32>
+  vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x8xi32>, memref<128x128xi32>
+  return
+}
+
+// -----
+
+//#########################################################
+// INT8 row-col-row
+//#########################################################
+
+// CHECK-DAG: [[$rowA0_map:#.+]] = affine_map<()[s0] -> (s0 mod 16)>
+// CHECK-DAG: [[$colA0_map:#.+]] = affine_map<()[s0] -> ((s0 floordiv 16) * 16)>
+// CHECK-DAG: [[$rowB0_map:#.+]] = affine_map<()[s0] -> (s0 mod 8)>
+// CHECK-DAG: [[$colB0_map:#.+]] = affine_map<()[s0] -> ((s0 floordiv 8) * 16)>
+
+// CHECK-DAG: [[$rowC0_map:#.+]] = affine_map<()[s0] -> (s0 floordiv 4)>
+// CHECK-DAG: [[$colC0_map:#.+]] = affine_map<()[s0] -> (s0 * 2 - (s0 floordiv 4) * 8)>
+// CHECK-DAG: [[$rowC8_map:#.+]] = affine_map<()[s0] -> (s0 floordiv 4 + 8)>
+
+
+#map0 = affine_map<(d0, d1) -> (d1, d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: func @m16n8k32_int8_row_col_row
+func.func @m16n8k32_int8_row_col_row(%arg0: memref<128x128xi8, 3>, %arg1: memref<128x128xi8, 3>, %arg2: memref<128x128xi32>) {
+  %cst_0 = arith.constant dense<0> : vector<32x8xi8>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %cst = arith.constant 0 : i8
+  %cst0 = arith.constant 0 : i32
+
+  // CHECK: [[lane:%.+]] = gpu.lane_id
+  // CHECK: [[row:%.+]] = affine.apply [[$rowA0_map]]()[[[lane]]]
+  // CHECK: [[col:%.+]] = affine.apply [[$colA0_map]]()[[[lane]]]
+  // CHECK: nvgpu.ldmatrix %arg0[[[row]], [[col]]] {numTiles = 4 : i32, transpose = false} : memref<128x128xi8, 3> -> vector<4x4xi8>
+
+  // CHECK: [[lane:%.+]] = gpu.lane_id
+  // CHECK: [[row:%.+]] = affine.apply [[$rowB0_map]]()[[[lane]]]
+  // CHECK: [[col:%.+]] = affine.apply [[$colB0_map]]()[[[lane]]]
+  // CHECK: nvgpu.ldmatrix %arg1[[[row]], [[col]]] {numTiles = 2 : i32, transpose = false} : memref<128x128xi8, 3> -> vector<2x4xi8>
+  
+  // CHECK: [[lane:%.+]] = gpu.lane_id
+  // CHECK: [[row:%.+]] = affine.apply [[$rowC0_map]]()[[[lane]]]
+  // CHECK: [[col:%.+]] = affine.apply [[$colC0_map]]()[[[lane]]]
+  // CHECK: vector.load %arg2[[[row]], [[col]]] : memref<128x128xi32>, vector<2xi32>
+  // CHECK: [[row:%.+]] = affine.apply [[$rowC8_map]]()[[[lane]]]
+  // CHECK: [[col:%.+]] = affine.apply [[$colC0_map]]()[[[lane]]]
+  // CHECK: vector.load %arg2[[[row]], [[col]]] : memref<128x128xi32>, vector<2xi32>
+  // CHECK-NOT: vector.load %arg2{{.*}}
+
+  %A = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : memref<128x128xi8, 3>, vector<16x32xi8>
+  %B = vector.transfer_read %arg1[%c0, %c0], %cst {in_bounds = [true, true]} : memref<128x128xi8, 3>, vector<8x32xi8>
+  %C = vector.transfer_read %arg2[%c0, %c0], %cst0 {in_bounds = [true, true]} : memref<128x128xi32>, vector<16x8xi32>
+  // CHECK: [[d:%.+]] = nvgpu.mma.sync({{.*}}) {mmaShape = [16, 8, 32]} : (vector<4x4xi8>, vector<2x4xi8>, vector<2x2xi32>) -> vector<2x2xi32>
+  %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %A, %B, %C : vector<16x32xi8>, vector<8x32xi8> into vector<16x8xi32>
+
+  // CHECK: [[lane:%.+]] = gpu.lane_id
+  // CHECK: [[v:%.+]] = vector.extract [[d]][0] : vector<2x2xi32>  
+  // CHECK: [[row:%.+]] = affine.apply [[$rowC0_map]]()[[[lane]]]
+  // CHECK: [[col:%.+]] = affine.apply [[$colC0_map]]()[[[lane]]]
+  // CHECK: vector.store [[v]], %arg2[[[row]], [[col]]] : memref<128x128xi32>, vector<2xi32>
+  // CHECK: [[v:%.+]] = vector.extract [[d]][1] : vector<2x2xi32>
+  // CHECK: [[row:%.+]] = affine.apply [[$rowC8_map]]()[[[lane]]]
+  // CHECK: [[col:%.+]] = affine.apply [[$colC0_map]]()[[[lane]]]
+  // CHECK: vector.store [[v]], %arg2[[[row]], [[col]]] : memref<128x128xi32>, vector<2xi32>
+  vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x8xi32>, memref<128x128xi32>
+  return
+}