[mlir][vectorToGPU] Fix type used when folding transpose into read op

[ThomasRaoux]

Summary

Pick the right result type when folding transpose op into a read

[qedawkins]

LGTM

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

440

What is this test testing differently from the one immediately above? I only notice the differing shapes.

[ThomasRaoux]

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

440

Yes, only the shapes are different, that's what exposes the bug since the pattern was mixing source and destination types.

[antiagainst]

IMO, the accessor name could use some improvements---it's very confusing and error-prone just being getVectorType(); making it something like getSrcVectorType() would make it far easier to avoid pitfalls. At least for me, I need to read the impl to double check what it means by just getVectorType().

[ThomasRaoux]

In D144113#4129557, @antiagainst wrote:

IMO, the accessor name could use some improvements---it's very confusing and error-prone just being getVectorType(); making it something like getSrcVectorType() would make it far easier to avoid pitfalls. At least for me, I need to read the impl to double check what it means by just getVectorType().

That's a good point. I agree.

[qedawkins]

In D144113#4129557, @antiagainst wrote:

IMO, the accessor name could use some improvements---it's very confusing and error-prone just being getVectorType(); making it something like getSrcVectorType() would make it far easier to avoid pitfalls. At least for me, I need to read the impl to double check what it means by just getVectorType().

+1 getVectorType() isn't by itself informative of its function, although with careful thought I could have avoided introducing this error. Also in this case it seems that none of the conversion tests were testing non-square matrices which was also a point of brittleness.

[antiagainst]

To close the loop: https://reviews.llvm.org/D144159 landed for improving the accessor methods.

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

 struct CombineTransferReadOpTranspose final
     : public OpRewritePattern<vector::TransposeOp> {
   using OpRewritePattern<vector::TransposeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::TransposeOp op,
                                 PatternRewriter &rewriter) const override {
     // Look through integer extend ops.
     Value source = op.getVector();
-    auto resultType = op.getVectorType();
+    Type resultType = op.getType();
     Operation *extOp;
     if ((extOp = source.getDefiningOp<arith::ExtSIOp>()) ||
         (extOp = source.getDefiningOp<arith::ExtUIOp>())) {
       source = extOp->getOperand(0);
       resultType =
-          VectorType::get(resultType.getShape(),
+          VectorType::get(resultType.cast<VectorType>().getShape(),
                           source.getType().cast<VectorType>().getElementType());
     }
 
     auto transferReadOp = source.getDefiningOp<vector::TransferReadOp>();
     if (!transferReadOp)
       return rewriter.notifyMatchFailure(op, "no transfer read");
 
     // TODO: support 0-d corner case.

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

func.func @matmul_mixed_signedness_int8(%arg0: memref<16x16xi8>, %arg1: memref<16x16xi8>, %arg2: memref<16x16xi32>) {
   %Br = vector.transfer_read %arg1[%c0, %c0], %cst_i8 {permutation_map = #map0, in_bounds = [true, true]} : memref<16x16xi8>, vector<16x16xi8>
   %C = vector.transfer_read %arg2[%c0, %c0], %cst_i32 {in_bounds = [true, true]} : memref<16x16xi32>, vector<16x16xi32>
   %Ae = arith.extui %Ar : vector<16x16xi8> to vector<16x16xi32>
   %Be = arith.extsi %Br : vector<16x16xi8> to vector<16x16xi32>
   %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %Ae, %Be, %C : vector<16x16xi32>, vector<16x16xi32> into vector<16x16xi32>
   vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xi32>, memref<16x16xi32>
   return
 }
+
+// -----
+
+#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)>
+#map4 = affine_map<(d0) -> (d0, 0)>
+#map5 = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK-LABEL: func @matmul_mixed_signedness_int8
+//   CHECK-DAG:   %[[A:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 32 : index} : memref<16x32xi8> -> !gpu.mma_matrix<16x32xui8, "AOp">
+//   CHECK-DAG:   %[[B:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 32 : index} : memref<16x32xi8> -> !gpu.mma_matrix<32x16xsi8, "BOp">
+//   CHECK-DAG:   %[[C:.+]] = gpu.subgroup_mma_load_matrix %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : memref<16x16xi32> -> !gpu.mma_matrix<16x16xi32, "COp">
+//       CHECK:   %[[D:.+]] = gpu.subgroup_mma_compute %[[A]], %[[B]], %[[C]] : !gpu.mma_matrix<16x32xui8, "AOp">, !gpu.mma_matrix<32x16xsi8, "BOp"> -> !gpu.mma_matrix<16x16xi32, "COp">
+//       CHECK:   gpu.subgroup_mma_store_matrix %[[D]], %{{.*}}[%{{.*}}, %{{.*}}] {leadDimension = 16 : index} : !gpu.mma_matrix<16x16xi32, "COp">, memref<16x16xi32>
+func.func @matmul_mixed_signedness_int8(%arg0: memref<16x32xi8>, %arg1: memref<16x32xi8>, %arg2: memref<16x16xi32>) {
+  %cst_0 = arith.constant dense<0> : vector<16x16xi8>
+  %c0 = arith.constant 0 : index
+  %cst_i8 = arith.constant 0 : i8
+  %cst_i32 = arith.constant 0 : i32
+  %Ar = vector.transfer_read %arg0[%c0, %c0], %cst_i8 {in_bounds = [true, true]} : memref<16x32xi8>, vector<16x32xi8>
+  %Br = vector.transfer_read %arg1[%c0, %c0], %cst_i8 {permutation_map = #map0, in_bounds = [true, true]} : memref<16x32xi8>, vector<16x32xi8>
+  %C = vector.transfer_read %arg2[%c0, %c0], %cst_i32 {in_bounds = [true, true]} : memref<16x16xi32>, vector<16x16xi32>
+  %Ae = arith.extui %Ar : vector<16x32xi8> to vector<16x32xi32>
+  %Be = arith.extsi %Br : vector<16x32xi8> to vector<16x32xi32>
+  %D = vector.contract {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %Ae, %Be, %C : vector<16x32xi32>, vector<16x32xi32> into vector<16x16xi32>
+  vector.transfer_write %D, %arg2[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xi32>, memref<16x16xi32>
+  return
+}
 No newline at end of file