[Mlir-commits] [mlir] [MLIR][NVGPU] Add convert.fpext and convert.fptrunc Ops (PR #199700)

[Srinivasa Ravi]

================
@@ -699,6 +699,128 @@ LogicalResult RcpOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// NVGPU_CvtFPTruncOp
+//===----------------------------------------------------------------------===//
+
+static bool isShapedContainerType(Type t) {
+  return llvm::isa<VectorType, RankedTensorType, UnrankedTensorType>(t);
+}
+
+static LogicalResult verifyConversionShapes(Operation *op, Type inType,
+                                            Type outType) {
+  if (llvm::isa<UnrankedTensorType>(inType) ||
+      llvm::isa<UnrankedTensorType>(outType))
+    return op->emitOpError("unranked tensor types are not supported, got ")
+           << inType << " and " << outType;
+  bool srcIsShaped = isShapedContainerType(inType);
+  bool dstIsShaped = isShapedContainerType(outType);
+  if (srcIsShaped != dstIsShaped)
+    return op->emitOpError("input and output must both be scalars or both be "
+                           "vectors/tensors, got ")
+           << inType << " and " << outType;
+  if (srcIsShaped) {
+    auto srcShaped = llvm::cast<ShapedType>(inType);
+    auto dstShaped = llvm::cast<ShapedType>(outType);
+    if (srcShaped.getRank() == 0 || dstShaped.getRank() == 0)
+      return op->emitOpError("rank-0 shaped types are not supported, use "
+                             "scalar type instead");
+    if (srcShaped.getShape() != dstShaped.getShape())
+      return op->emitOpError("input and output shapes must match, got ")
+             << inType << " and " << outType;
+    if (llvm::isa<VectorType>(inType) != llvm::isa<VectorType>(outType))
+      return op->emitOpError("input and output must be the same container "
+                             "type (both vector or both tensor), got ")
+             << inType << " and " << outType;
+  }
+  return success();
+}
+
+LogicalResult FPTruncOp::verify() {
+  Type inType = getIn().getType();
+  Type outType = getType();
+  Type srcType = getElementTypeOrSelf(inType);
+  Type dstType = getElementTypeOrSelf(outType);
+  int srcBitWidth = srcType.getIntOrFloatBitWidth();
+  int dstBitWidth = dstType.getIntOrFloatBitWidth();
+  auto rnd = getRnd();
+
+  if (auto result = verifyConversionShapes(getOperation(), inType, outType);
+      failed(result))
+    return result;
+
+  if (srcBitWidth <= dstBitWidth)
+    return emitOpError("result type ")
+           << dstType << " must be narrower than operand type " << srcType;
+
+  if (!(srcBitWidth == 64 || srcBitWidth == 32 || srcBitWidth == 16))
+    return emitOpError("input type must be 64/32/16 bitwidth, but got ")
+           << srcBitWidth;
+
+  if (llvm::isa<Float8E8M0FNUType>(dstType)) {
+    if (rnd != mlir::NVVM::FPRoundingMode::RZ &&
+        rnd != mlir::NVVM::FPRoundingMode::RP)
+      return emitOpError("expects RZ or RP rounding mode when result type is "
+                         "e8m0, but got ")
+             << getRndAttr();
+  } else if (rnd == mlir::NVVM::FPRoundingMode::RS) {
+    // TODO: Currently, we only support conversions which fit into a single i32
+    // register. Support f32->f8/f6/f4 conversions with RS rounding.
+    if (!(srcBitWidth == 32 && (dstBitWidth == 16)))
+      return emitOpError("RS (stochastic) rounding is only supported for "
+                         "f32->f16/bf16, got ")
+             << srcType << " -> " << dstType;
+    if (!getRandomBits())
+      return emitOpError("random_bits operand is required with RS rounding");
+  } else if (rnd != mlir::NVVM::FPRoundingMode::RN) {
----------------
Wolfram70 wrote:

Agreed, added support for `RZ` in the latest revision. Thanks for pointing this out!

https://github.com/llvm/llvm-project/pull/199700