[mlir][ROCDL] Address review comments

- Use SmallVectorImpl for output parameters per LLVM coding standards
- Fix type casting bug (getSourceA -> getSourceB on line 1309)
- Use cast<> instead of dyn_cast<> for guaranteed vector types
- Simplify assembly format with functional-type shorthand

Addresses review comments from @kuhar on PR #167041

Author: Muzammiluddin-Syed-ECE

mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td

@@ -590,7 +590,7 @@ class ROCDL_WMMA_IntrOp<string mnemonic, Type AB, Type CD> : ROCDL_IntrOp<mnemon
              LLVM_ScalarOrVectorOf<CD>:$C)> {
   let results = (outs LLVM_ScalarOrVectorOf<CD>:$res);
   let assemblyFormat = [{
-    $A `,` $B `,` $C attr-dict `:` `(`type($A) `,` type($B) `,` type($C)`)` `->` type($res)
+    $A `,` $B `,` $C attr-dict `:` functional-type(operands, $res)
   }];
 }
 
@@ -603,7 +603,7 @@ class ROCDL_WMMA_Opsel_IntrOp<string mnemonic, Type AB, Type CD> : ROCDL_IntrOp<
              DefaultValuedAttr<I1Attr, "0">:$opsel)> {
   let results = (outs LLVM_ScalarOrVectorOf<CD>:$res);
   let assemblyFormat = [{
-    $A `,` $B `,` $C attr-dict `:` `(`type($A) `,` type($B) `,` type($C)`)` `->` type($res)
+    $A `,` $B `,` $C attr-dict `:` functional-type(operands, $res)
   }];
 }
 
@@ -618,7 +618,7 @@ class ROCDL_WMMA_IU_IntrOp<string mnemonic, Type AB, Type CD> : ROCDL_IntrOp<mne
              DefaultValuedAttr<I1Attr, "0">:$clamp)> {
   let results = (outs LLVM_ScalarOrVectorOf<CD>:$res);
   let assemblyFormat = [{
-    $A `,` $B `,` $C attr-dict `:` `(`type($A) `,` type($B) `,` type($C)`)` `->` type($res)
+    $A `,` $B `,` $C attr-dict `:` functional-type(operands, $res)
   }];
 }
 
@@ -635,7 +635,7 @@ class ROCDL_WMMA_ModsAll_Reuse_IntrOp<string mnemonic, Type AB, Type CD> : ROCDL
              DefaultValuedAttr<I1Attr, "0">:$reuseB)> {
   let results = (outs LLVM_ScalarOrVectorOf<CD>:$res);
   let assemblyFormat = [{
-    $A `,` $B `,` $C attr-dict `:` `(`type($A) `,` type($B) `,` type($C)`)` `->` type($res)
+    $A `,` $B `,` $C attr-dict `:` functional-type(operands, $res)
   }];
 }
 
@@ -650,7 +650,7 @@ class ROCDL_WMMA_ModsC_IntrOp<string mnemonic, Type AB, Type CD> : ROCDL_IntrOp<
              DefaultValuedAttr<I1Attr, "0">:$reuseB)> {
   let results = (outs LLVM_ScalarOrVectorOf<CD>:$res);
   let assemblyFormat = [{
-    $A `,` $B `,` $C attr-dict `:` `(`type($A) `,` type($B) `,` type($C)`)` `->` type($res)
+    $A `,` $B `,` $C attr-dict `:` functional-type(operands, $res)
   }];
 }
 
@@ -667,7 +667,7 @@ class ROCDL_WMMA_ModsAll_Diff_IntrOp<string mnemonic, Type AB, Type C, Type D> :
              DefaultValuedAttr<I1Attr, "0">:$reuseB)> {
   let results = (outs LLVM_ScalarOrVectorOf<D>:$res);
   let assemblyFormat = [{
-    $A `,` $B `,` $C attr-dict `:` `(`type($A) `,` type($B) `,` type($C)`)` `->` type($res)
+    $A `,` $B `,` $C attr-dict `:` functional-type(operands, $res)
   }];
 }
 
@@ -683,7 +683,7 @@ class ROCDL_WMMA_ModsAB_IntrOp<string mnemonic, Type AB, Type CD> : ROCDL_IntrOp
              DefaultValuedAttr<I1Attr, "0">:$reuseB)> {
   let results = (outs LLVM_ScalarOrVectorOf<CD>:$res);
   let assemblyFormat = [{
-    $A `,` $B `,` $C attr-dict `:` `(`type($A) `,` type($B) `,` type($C)`)` `->` type($res)
+    $A `,` $B `,` $C attr-dict `:` functional-type(operands, $res)
   }];
 }
 

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

@@ -682,8 +682,8 @@ static Value castMFMAScaleOperand(ConversionPatternRewriter &rewriter,
 static void wmmaPushInputOperand(
     ConversionPatternRewriter &rewriter, Location loc,
     const TypeConverter *typeConverter, bool isUnsigned, Value llvmInput,
-    Value mlirInput, SmallVector<Value, 4> &operands,
-    SmallVector<NamedAttribute, 4> &attrs, StringRef attrName) {
+    Value mlirInput, SmallVectorImpl<Value> &operands,
+    SmallVectorImpl<NamedAttribute> &attrs, StringRef attrName) {
   Type inputType = llvmInput.getType();
   auto vectorType = dyn_cast<VectorType>(inputType);
   if (!vectorType) {
@@ -741,8 +741,8 @@ static void wmmaPushOutputOperand(ConversionPatternRewriter &rewriter,
                                   Location loc,
                                   const TypeConverter *typeConverter,
                                   Value output, int32_t subwordOffset,
-                                  bool clamp, SmallVector<Value, 4> &operands,
-                                  SmallVector<NamedAttribute, 4> &attrs) {
+                                  bool clamp, SmallVectorImpl<Value> &operands,
+                                  SmallVectorImpl<NamedAttribute> &attrs) {
   Type inputType = output.getType();
   auto vectorType = dyn_cast<VectorType>(inputType);
   Type elemType = vectorType.getElementType();
@@ -1305,9 +1305,9 @@ struct WMMAOpLowering : public ConvertOpToLLVMPattern<WMMAOp> {
     // The WMMA operations represent vectors of bf16s as vectors of i16s
     // (except on gfx1250), so we need to bitcast bfloats to i16 and then
     // bitcast them back.
-    auto aType = dyn_cast<VectorType>(adaptor.getSourceA().getType());
-    auto bType = dyn_cast<VectorType>(adaptor.getSourceA().getType());
-    auto destCType = dyn_cast<VectorType>(adaptor.getDestC().getType());
+    auto aType = cast<VectorType>(adaptor.getSourceA().getType());
+    auto bType = cast<VectorType>(adaptor.getSourceB().getType());
+    auto destCType = cast<VectorType>(adaptor.getDestC().getType());
     bool castAToI16 = aType.getElementType().isBF16() && !isGFX1250;
     bool castBToI16 = bType.getElementType().isBF16() && !isGFX1250;
     bool castDestCToI16 = destCType.getElementType().isBF16() && !isGFX1250;