[mlir][nvpu] Prevent F32ToTF32 pattern to generate illegal IR

[ThomasRaoux]

Summary

We shouldn't apply this pattern to non F32->F32 mma.sync operations.

mlir/lib/Dialect/NVGPU/Transforms/MmaSyncTF32Transform.cpp

MmaSyncF32ToTF32Pattern(MLIRContext *context,
                           nvgpu::MmaSyncF32Lowering precision)
       : OpRewritePattern<nvgpu::MmaSyncOp>(context, /*benifit*/ 1),
         precision(precision) {}
 
   LogicalResult matchAndRewrite(nvgpu::MmaSyncOp op,
                                 PatternRewriter &rewrite) const override {
     Location location = op->getLoc();
 
-    if (op->hasAttr(op.getTf32EnabledAttrName()))
+    if (op->hasAttr(op.getTf32EnabledAttrName()) ||
+        !op.getMatrixA().getType().cast<VectorType>().getElementType().isF32())
       return failure();
 
     if (precision == MmaSyncF32Lowering::Unkown)
       return emitError(location, "MmaSync F32-to-TF32 cannot be lowered with "
                                  "unknown precision level");
 
     if (precision == MmaSyncF32Lowering::TF32x3)
       return emitError(location, "TF32x3 is not supported at the moment "

mlir/test/Dialect/NVGPU/mma-sync-f32-to-tf32.mlir

 // CHECK-LABEL: m16n8k8_tf32
 func.func @m16n8k8_tf32(%arg0: vector<4x1xf32>, %arg1: vector<2x1xf32>, %arg2: vector<2x2xf32>) -> vector<2x2xf32> {
   // CHECK: nvgpu.mma.sync
   // CHECK-SAME: tf32Enabled
   %d = nvgpu.mma.sync (%arg0, %arg1, %arg2) {mmaShape = [16, 8, 8]} : (vector<4x1xf32>, vector<2x1xf32>, vector<2x2xf32>) -> vector<2x2xf32>    
   return %d : vector<2x2xf32>
 }
 // -----
+
+// Negative test for non f32 case.
+// CHECK-LABEL: mma_sync_f16
+//   CHECK-NOT: tf32Enabled
+//       CHECK: return
+func.func @mma_sync_f16(%arg0: vector<4x2xf16>, %arg1: vector<2x2xf16>, %arg2: vector<2x2xf16>) -> vector<2x2xf16> {
+  // expected-error @+1 {{TF32x3 is not supported at the moment for nvgpu.mma.sync on f32 datatype}}
+  %d = nvgpu.mma.sync (%arg0, %arg1, %arg2) {mmaShape = [16, 8, 16]} : (vector<4x2xf16>, vector<2x2xf16>, vector<2x2xf16>) -> vector<2x2xf16>
+  return %d : vector<2x2xf16>
+}