Enable BF16 casts for non-SPIRV targets. (#3992)

Signed-off-by: Ilya Enkovich <[email protected]>

Author: ienkovich

third_party/intel/lib/TritonIntelGPUToLLVM/BF16Casts.cpp

@@ -10,6 +10,7 @@
 #include "llvm/ADT/TypeSwitch.h"
 
 #include "intel/include/Dialect/TritonIntelGPU/IR/Dialect.h"
+#include "intel/include/Dialect/TritonIntelGPU/IR/Utils.h"
 
 using namespace mlir;
 
@@ -83,23 +84,29 @@ Value convertBf16ToFp32(Location loc, ConversionPatternRewriter &rewriter,
     auto moduleOp = definingOp->getParentWithTrait<OpTrait::SymbolTable>();
     if (moduleOp->hasAttr(triton::gpu::intel::TritonIntelGPUDialect::
                               getSupportBF16ConversionAttrName())) {
-      constexpr StringLiteral baseName = "__spirv_ConvertBF16ToFINTEL";
-      Type inTy = getTypeWithSameShape(v.getType(), i16_ty);
-      Type outTy = getTypeWithSameShape(inTy, f32_ty);
-      std::string funcName = mlir::triton::gpu::intel::mangle(baseName, inTy);
-
-      auto bitcastValue = b.bitcast(v, inTy).getResult();
-
-      auto memAttr = rewriter.getAttr<LLVM::MemoryEffectsAttr>(
-          /*other=*/LLVM::ModRefInfo::NoModRef,
-          /*argMem=*/LLVM::ModRefInfo::NoModRef,
-          /*inaccessibleMem=*/LLVM::ModRefInfo::NoModRef);
-      auto funcAttrs = gpu::intel::noUnwindWillReturnAttrs;
-      funcAttrs.memEffectsAttr = memAttr;
-
-      auto call = gpu::intel::createDeviceFunctionCall(
-          rewriter, funcName, outTy, {inTy}, {bitcastValue}, {}, funcAttrs);
-      return call.getResult();
+      // For SPIRV target, use specialized intrinsic call for conversion.
+      // Otherwise, use fpext operation.
+      if (gpu::intel::hasSpirvTargetArch(moduleOp)) {
+        constexpr StringLiteral baseName = "__spirv_ConvertBF16ToFINTEL";
+        Type inTy = getTypeWithSameShape(v.getType(), i16_ty);
+        Type outTy = getTypeWithSameShape(inTy, f32_ty);
+        std::string funcName = mlir::triton::gpu::intel::mangle(baseName, inTy);
+
+        auto bitcastValue = b.bitcast(v, inTy).getResult();
+
+        auto memAttr = rewriter.getAttr<LLVM::MemoryEffectsAttr>(
+            /*other=*/LLVM::ModRefInfo::NoModRef,
+            /*argMem=*/LLVM::ModRefInfo::NoModRef,
+            /*inaccessibleMem=*/LLVM::ModRefInfo::NoModRef);
+        auto funcAttrs = gpu::intel::noUnwindWillReturnAttrs;
+        funcAttrs.memEffectsAttr = memAttr;
+
+        auto call = gpu::intel::createDeviceFunctionCall(
+            rewriter, funcName, outTy, {inTy}, {bitcastValue}, {}, funcAttrs);
+        return call.getResult();
+      }
+
+      return rewriter.create<LLVM::FPExtOp>(loc, f32_ty, v);
     }
   }
 
@@ -118,22 +125,27 @@ Value convertFp32ToBf16(Location loc, ConversionPatternRewriter &rewriter,
                               getSupportBF16ConversionAttrName()) &&
         rounding == RoundingMode::RTNE) {
       // Intel SPIR-V extension only supports round-to-nearest-even
-      constexpr StringLiteral baseName = "__spirv_ConvertFToBF16INTEL";
-      Type inTy = v.getType();
-      Type funcOutTy = getTypeWithSameShape(inTy, i16_ty);
-      Type outTy = getTypeWithSameShape(inTy, bf16_ty);
-      std::string funcName = mlir::triton::gpu::intel::mangle(baseName, inTy);
-
-      auto memAttr = rewriter.getAttr<LLVM::MemoryEffectsAttr>(
-          /*other=*/LLVM::ModRefInfo::NoModRef,
-          /*argMem=*/LLVM::ModRefInfo::NoModRef,
-          /*inaccessibleMem=*/LLVM::ModRefInfo::NoModRef);
-      auto funcAttrs = gpu::intel::noUnwindWillReturnAttrs;
-      funcAttrs.memEffectsAttr = memAttr;
-
-      auto call = gpu::intel::createDeviceFunctionCall(
-          rewriter, funcName, funcOutTy, {inTy}, {v}, {}, funcAttrs);
-      return b.bitcast(call.getResult(), outTy);
+      // LLVM fptrunc operation also assumes round-to-nearest mode
+      if (gpu::intel::hasSpirvTargetArch(moduleOp)) {
+        constexpr StringLiteral baseName = "__spirv_ConvertFToBF16INTEL";
+        Type inTy = v.getType();
+        Type funcOutTy = getTypeWithSameShape(inTy, i16_ty);
+        Type outTy = getTypeWithSameShape(inTy, bf16_ty);
+        std::string funcName = mlir::triton::gpu::intel::mangle(baseName, inTy);
+
+        auto memAttr = rewriter.getAttr<LLVM::MemoryEffectsAttr>(
+            /*other=*/LLVM::ModRefInfo::NoModRef,
+            /*argMem=*/LLVM::ModRefInfo::NoModRef,
+            /*inaccessibleMem=*/LLVM::ModRefInfo::NoModRef);
+        auto funcAttrs = gpu::intel::noUnwindWillReturnAttrs;
+        funcAttrs.memEffectsAttr = memAttr;
+
+        auto call = gpu::intel::createDeviceFunctionCall(
+            rewriter, funcName, funcOutTy, {inTy}, {v}, {}, funcAttrs);
+        return b.bitcast(call.getResult(), outTy);
+      }
+
+      return rewriter.create<LLVM::FPTruncOp>(loc, bf16_ty, v);
     }
   }