Add Fp4 lowering from MIGraphX Dialect (#2089)

Add lowering for Fp4 scaled GEMM in migraphx dialect

Author: umangyadav

external/llvm-project/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp

@@ -1254,7 +1254,8 @@ LogicalResult TosaValidation::applyErrorIfCheck(Operation *op) {
 bool TosaValidation::isValidElementType(Type type, const bool allowUnsigned) {
   if (isa<FloatType>(type)) {
     return isa<Float32Type, Float16Type, BFloat16Type, Float8E4M3FNType,
-               Float8E5M2Type, Float8E4M3FNUZType, Float8E5M2FNUZType>(type);
+               Float8E5M2Type, Float8E4M3FNUZType, Float8E5M2FNUZType,
+               Float4E2M1FNType, Float8E8M0FNUType>(type);
   }
   if (auto intTy = dyn_cast<IntegerType>(type)) {
     if (intTy.isSignless()) {

mlir/include/mlir/Dialect/MIGraphX/IR/MIGraphX.td

@@ -523,31 +523,47 @@ def MIGraphX_MultiBroadcastOp :
   }];
 }
 
-class MIGraphX_DotOpBase<string mnemonic, list<Type> inputTypes=[], list<Type> outputTypes=[]> :
-    MIGraphX_Op<mnemonic>,
-    Arguments<(ins MIXRShapedOf<inputTypes>:$in_a,
-                   MIXRShapedOf<inputTypes>:$in_b
-                   )>,
-	Results<(outs MIXRShapedOf<outputTypes>:$output)> {
-    let assemblyFormat = [{
-      $in_a `,` $in_b attr-dict `:` type($in_a) `,` type($in_b) `->` type($output)
-    }];
-}
-
-def MIGraphX_QuantDotOp :
-    MIGraphX_DotOpBase<"quant_dot", [F8E4M3FNUZ, F8E5M2FNUZ, F8E4M3FN, F8E5M2, I8, SI8], [F32, I32, SI32]>{
+defvar QuantDotInTypes = [F8E4M3FNUZ, F8E5M2FNUZ, F8E4M3FN, F8E5M2, F4E2M1FN,
+                          I8, SI8];
+defvar QuantDotOutTypes = [F32, I32, SI32];
+defvar QuantDotScaleTypes = [F8E8M0FNU, F32];
+
+def MIGraphX_QuantDotOp
+    : MIGraphX_Op<"quant_dot", [AttrSizedOperandSegments]>,
+      Arguments<(ins MIXRShapedOf<QuantDotInTypes>:$in_a,
+          MIXRShapedOf<QuantDotInTypes>:$in_b,
+          Optional<MIXRShapedOf<QuantDotScaleTypes>>:$scaleA,
+          Optional<MIXRShapedOf<QuantDotScaleTypes>>:$scaleB)>,
+      Results<(outs MIXRShapedOf<QuantDotOutTypes>:$output)> {
   let summary = "Dot product of quantized tensors";
   let description = [{
-    The `migraphx.quant_dot` op computes the dot product of two tensors.
+    The `migraphx.quant_dot` op computes the dot product of two tensors. This operation is used when converting MIGraphX IR's quant_dot operation to MIGraphX Dialect. 
+    Usually `migraphx.quant_dot` operations have QuantizeLinear and DequantizeLinear operations to convert data types from higher precision to lower precision.  
+    `migraphx.quant_dot` operation is also used for "scaled" GEMMs. 
   }];
+  let assemblyFormat = [{
+    $in_a (`scaled` `by` $scaleA^)? `,` $in_b (`scaled` `by` $scaleB^)? attr-dict
+    `:` type($in_a) (`scaled` `by` type($scaleA)^)? `,` type($in_b)
+        (`scaled` `by` type($scaleB)^)?
+    `->` type($output)
+  }];
+  let hasVerifier = 1;
 }
 
-def MIGraphX_DotOp :
-    MIGraphX_DotOpBase<"dot", [F32, F16, BF16], [F32, F16, BF16]>{
+defvar DotInTypes = [F32, F16, BF16, F4E2M1FN];
+defvar DotOutTypes = [F32, F16, BF16];
+
+def MIGraphX_DotOp : MIGraphX_Op<"dot">,
+                     Arguments<(ins MIXRShapedOf<DotInTypes>:$in_a,
+                         MIXRShapedOf<DotInTypes>:$in_b)>,
+                     Results<(outs MIXRShapedOf<DotOutTypes>:$output)> {
   let summary = "Dot product of tensors";
   let description = [{
     The `migraphx.dot` op computes the dot product of two tensors.
   }];
+  let assemblyFormat = [{
+    $in_a `,` $in_b attr-dict `:` type($in_a) `,` type($in_b) `->` type($output)
+  }];
 }
 
 def MIGraphX_SoftmaxOp :