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

[Guray Ozen]

================
@@ -1709,6 +1711,685 @@ struct NVGPURcpOpLowering : public ConvertOpToLLVMPattern<nvgpu::RcpOp> {
         rewriter);
   }
 };
+
+//===----------------------------------------------------------------------===//
+// NVGPUConvertFPTruncOp Lowering
+//===----------------------------------------------------------------------===//
+
+enum class FPKind { F32, BF16, F16, F8, F6, F4 };
+
+static int getEffectiveBitWidth(int bitWidth) {
+  // f6 types are 6-bit but NVVM Ops expect 8-bit (i8) containers.
+  return bitWidth == 6 ? 8 : bitWidth;
+}
+
+static std::optional<FPKind> classifyFPType(Type t) {
+  static constexpr auto isConvertibleF8Type = [](Type t) {
+    return isa<Float8E4M3FNType, Float8E5M2Type, Float8E8M0FNUType>(t);
+  };
+  static constexpr auto isConvertibleF6Type = [](Type t) {
+    return isa<Float6E2M3FNType, Float6E3M2FNType>(t);
+  };
+  static constexpr auto isConvertibleF4Type = [](Type t) {
+    return isa<Float4E2M1FNType>(t);
+  };
+
+  if (t.isF32())
+    return FPKind::F32;
+  if (t.isBF16())
+    return FPKind::BF16;
+  if (t.isF16())
+    return FPKind::F16;
+  if (isConvertibleF8Type(t))
+    return FPKind::F8;
+  if (isConvertibleF6Type(t))
+    return FPKind::F6;
+  if (isConvertibleF4Type(t))
+    return FPKind::F4;
+  return std::nullopt;
+}
+
+/// Number of source-side i32 register slots consumed by each NVVM convert Op.
+static int getNumSrcI32PerConv(FPKind src) {
+  return src == FPKind::F32 ? 2 : 1;
+}
+
+/// Conversion op identifier for nvgpu.convert.fptrunc lowering dispatch table.
+enum class FPTruncConvOp {
+  F32x2_TO_F16x2,
+  F32x2_TO_BF16x2,
+  F32x2_TO_F8x2,
+  F32x2_TO_F6x2,
+  F32x2_TO_F4x2,
+  F16x2_TO_F8x2,
+  F16x2_TO_F6x2,
+  F16x2_TO_F4x2,
+  BF16x2_TO_F8x2,
+  BF16x2_TO_F6x2,
+  BF16x2_TO_F4x2,
+};
+
+struct FPTruncTableEntry {
+  FPKind src;
+  FPKind dst;
+  FPTruncConvOp convOp;
+};
+
+static constexpr FPTruncTableEntry kFPTruncTable[] = {
+    // f32 source
+    {FPKind::F32, FPKind::F16, FPTruncConvOp::F32x2_TO_F16x2},
+    {FPKind::F32, FPKind::BF16, FPTruncConvOp::F32x2_TO_BF16x2},
+    {FPKind::F32, FPKind::F8, FPTruncConvOp::F32x2_TO_F8x2},
+    {FPKind::F32, FPKind::F6, FPTruncConvOp::F32x2_TO_F6x2},
+    {FPKind::F32, FPKind::F4, FPTruncConvOp::F32x2_TO_F4x2},
+    // f16 source
+    {FPKind::F16, FPKind::F8, FPTruncConvOp::F16x2_TO_F8x2},
+    {FPKind::F16, FPKind::F6, FPTruncConvOp::F16x2_TO_F6x2},
+    {FPKind::F16, FPKind::F4, FPTruncConvOp::F16x2_TO_F4x2},
+    // bf16 source
+    {FPKind::BF16, FPKind::F8, FPTruncConvOp::BF16x2_TO_F8x2},
+    {FPKind::BF16, FPKind::F6, FPTruncConvOp::BF16x2_TO_F6x2},
+    {FPKind::BF16, FPKind::F4, FPTruncConvOp::BF16x2_TO_F4x2},
+};
+
+static std::optional<FPTruncTableEntry> lookupTruncConvOp(Type srcElemType,
+                                                          Type dstElemType) {
+  auto srcKind = classifyFPType(srcElemType);
----------------
grypp wrote:

please don't use auto when the type is not obvious

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