Add f32 RTNE to tf32 in DPAS (#3803)

SPIRV extension `_Z25__spirv_RoundFToTF32INTELf` for fp32 to tf32.

---------

Co-authored-by: Si, Yudong <yudong.si@intel.com>
Co-authored-by: Whitney Tsang <whitney.tsang@intel.com>

Author: 3 people

test/Conversion/intel/tritongpu_to_gen_dot.mlir

@@ -72,9 +72,20 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
 #dot_operand_b = #ttg.dot_op<{opIdx=1, parent=#dpas, kWidth=1}>
 
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
-  // CHECK-LABEL: dot_f32_tf32_tf32_f32_1
+  // CHECK-LABEL: llvm.func spir_kernelcc @dot_f32_tf32_tf32_f32_1(
+  // CHECK-SAME:    %[[A:.*]]: !llvm.struct<(f32, f32, f32, f32)>, %[[B:.*]]: !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32)>,
+  // CHECK-SAME:    %[[C:.*]]: !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32)>) attributes {intel_reqd_sub_group_size = 32 : i32, triton_gen.max_work_group_size = array<i32: 32, 1, 1>} {
   tt.func @dot_f32_tf32_tf32_f32_1(%a: tensor<8x8xf32, #dot_operand_a>, %b: tensor<8x16xf32, #dot_operand_b>, %c: tensor<8x16xf32, #dpas>) {
-    // CHECK: llvm.call spir_funccc @_Z39intel_sub_group_tf32_tf32_matrix_mad_k8Dv4_fDv8_fS0_(%{{.*}}, %{{.*}}, %{{.*}}) {{.*}} : (vector<4xf32>, vector<8xf32>, vector<8xf32>) -> vector<8xf32>
+    // COM: To simplify, only check RTNE and its usage for the last element of A, B, C
+    // CHECK %[[A_LAST_VAL:.*]] = llvm.extractvalue %[[A]][3]
+    // CHECK %[[A_RTNE_VAL:.*]] = llvm.call spir_funccc @_Z25__spirv_RoundFToTF32INTELf(%[[A_LAST_VAL]])
+    // CHECK %[[A_0:.*]] = llvm.insertelement %[[A_RTNE_VAL]], %{{.*}}{{\[}}%{{.*}} : i32] : vector<4xf32>
+    // CHECK %[[B_LAST_VAL:.*]] = llvm.extractvalue %[[B]][7]
+    // CHECK %[[B_RTNE_VAL:.*]] = llvm.call spir_funccc @_Z25__spirv_RoundFToTF32INTELf(%[[B_LAST_VAL]])
+    // CHECK %[[B_0:.*]] = llvm.insertelement %[[B_RTNE_VAL]], %{{.*}}{{\[}}%{{.*}} : i32] : vector<8xf32>
+    // CHECK %[[C_LAST_VAL:.*]] = llvm.extractvalue %[[C]][7]
+    // CHECK %[[C_0:.*]] = llvm.insertelement %[[C_LAST_VAL]], %{{.*}}{{\[}}%{{.*}} : i32] : vector<8xf32>
+    // CHECK : llvm.call spir_funccc @_Z39intel_sub_group_tf32_tf32_matrix_mad_k8Dv4_fDv8_fS0_(%[[A_0]], %[[B_0]], %[[C_0]]) {{.*}} : (vector<4xf32>, vector<8xf32>, vector<8xf32>) -> vector<8xf32>
     %0 = tt.dot %a, %b, %c, inputPrecision = tf32 : tensor<8x8xf32, #dot_operand_a> * tensor<8x16xf32, #dot_operand_b> -> tensor<8x16xf32, #dpas>
     tt.return
   }

third_party/intel/include/Dialect/TritonGEN/IR/TritonGENOps.td

@@ -20,6 +20,7 @@ include "mlir/IR/EnumAttr.td"
 include "mlir/Dialect/LLVMIR/LLVMTypes.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
+include "mlir/Interfaces/InferTypeOpInterface.td" // SameOperandsAndResultType
 
 //===----------------------------------------------------------------------===//
 // TritonGEN op definitions
@@ -313,4 +314,19 @@ def TritonGEN_SubGroupBlockWriteOp : TritonGEN_Op<"sub_group_block_write"> {
   }];
 }
 
+def TritonGEN_FToTf32Op
+    : TritonGEN_Op<"f_to_tf32", [SameOperandsAndResultType]> {
+  let summary = "Rounding instruction from float to tensor float (TF32) data format";
+
+  let description = [{
+    The op converts value numerically from
+    a 32-bit floating point type to TF32 with rounding to the nearest even.
+  }];
+
+  let arguments = (ins F32:$val);
+  let results = (outs F32:$res);
+  let assemblyFormat = [{
+    $val attr-dict `:` type($val)
+  }];
+}
 #endif // TRITONGEN_OPS

third_party/intel/lib/Target/SPIRV/SPIRVTranslation.cpp

@@ -107,7 +107,7 @@ class SmallVectorBuffer : public std::streambuf {
 
 static SPIRV::TranslatorOpts getSPIRVOopts() {
   SPIRV::TranslatorOpts SPIRVOpts;
-  static constexpr std::array<SPIRV::ExtensionID, 12> AllowedExtensions{
+  static constexpr std::array<SPIRV::ExtensionID, 13> AllowedExtensions{
       SPIRV::ExtensionID::SPV_EXT_shader_atomic_float_add,
       SPIRV::ExtensionID::SPV_INTEL_arbitrary_precision_integers,
       SPIRV::ExtensionID::SPV_INTEL_arithmetic_fence,
@@ -116,6 +116,7 @@ static SPIRV::TranslatorOpts getSPIRVOopts() {
       SPIRV::ExtensionID::SPV_INTEL_kernel_attributes,
       SPIRV::ExtensionID::SPV_INTEL_memory_access_aliasing,
       SPIRV::ExtensionID::SPV_INTEL_subgroups,
+      SPIRV::ExtensionID::SPV_INTEL_tensor_float32_rounding,
       SPIRV::ExtensionID::SPV_INTEL_unstructured_loop_controls,
       SPIRV::ExtensionID::SPV_INTEL_vector_compute,
       SPIRV::ExtensionID::SPV_KHR_bit_instructions,

third_party/intel/lib/TritonGENToLLVM/TritonGENToLLVMPass.cpp

@@ -699,6 +699,31 @@ struct TritonSubGroupBlockWriteLowering
   }
 };
 
+struct TritonFToTf32OpLowering
+    : public ConvertOpToLLVMPattern<TritonGEN::FToTf32Op> {
+  using ConvertOpToLLVMPattern<TritonGEN::FToTf32Op>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(TritonGEN::FToTf32Op op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    MLIRContext *ctx = rewriter.getContext();
+    Location loc = op->getLoc();
+    auto b = TritonLLVMOpBuilder(loc, rewriter);
+
+    Value value = op->getOperand(0);
+    SmallVector<Type> argTypes{f32_ty};
+    SmallVector<Value> args{value};
+
+    const StringLiteral funcName = "_Z25__spirv_RoundFToTF32INTELf";
+    auto retType = f32_ty;
+    auto callOp = intel::createDeviceFunctionCall(
+        rewriter, funcName, retType, {argTypes}, {args}, {},
+        intel::noUnwindWillReturnAttrs);
+    rewriter.replaceOp(op, callOp);
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
@@ -760,7 +785,8 @@ void mlir::triton::populateTritonGENToLLVMConversionPatterns(
       .add<TritonMatrixDPASLowering, TritonMatrix2DBlockLoadLowering,
            TritonMatrix2DBlockStoreLowering,
            TritonMatrix2DBlockPrefetchLowering, TritonSubGroupBlockReadLowering,
-           TritonSubGroupBlockWriteLowering>(converter);
+           TritonSubGroupBlockWriteLowering, TritonFToTf32OpLowering>(
+          converter);
 }
 
 void registerConvertTritonTritonGENToLLVMInterface(DialectRegistry &registry) {

third_party/intel/lib/TritonIntelGPUToLLVM/DotOpToLLVM/DPAS.cpp

@@ -310,16 +310,21 @@ class DotOpDPASConversionHelper {
     size_t rank = repCluster.size();
     unsigned repClusterOuter = 0u;
     unsigned repClusterInner = 0u;
+    bool isOperandA = false;
+    bool isOperandB = false;
+    bool isFToTF32Enabled = false;
     switch (opIdx) {
     case DpasEncodingAttr::OpIdx::OperandA:
       // operand A
       repClusterOuter = repCluster[rank - 2];
       repClusterInner = 1;
+      isOperandA = true;
       break;
     case DpasEncodingAttr::OpIdx::OperandB:
       // operand B
       repClusterInner = 1;
       repClusterOuter = repCluster[rank - 1];
+      isOperandB = true;
       break;
     case DpasEncodingAttr::OpIdx::OperandC:
       // operand C
@@ -333,6 +338,11 @@ class DotOpDPASConversionHelper {
         totalElems /
         ((batch * outer * inner) * (repClusterOuter * repClusterInner));
     VectorType dotOpTy = vec_ty(elemTy, numElemsPerOperand);
+    // TODO: IGC bug, Update isFToTF32Enabled as follows once issue #3870 is
+    // fixed. isFToTF32Enabled = elemTy.isFloat(32) && (isOperandA ||
+    // isOperandB)
+    isFToTF32Enabled = elemTy.isFloat(32) &&
+                       ((rank == 3) ? isOperandA : (isOperandA || isOperandB));
 
     auto tb = TritonLLVMOpBuilder(loc, rewriter);
     int offset = 0;
@@ -344,8 +354,18 @@ class DotOpDPASConversionHelper {
             for (int repInner = 0; repInner < repClusterInner; ++repInner) {
               Value matVal = rewriter.create<LLVM::UndefOp>(loc, dotOpTy);
               for (int k = 0; k < numElemsPerOperand; ++k) {
-                matVal = tb.insert_element(dotOpTy, matVal, elems[offset++],
-                                           tb.i32_val(k));
+                if (isFToTF32Enabled) {
+                  Value f32Val = elems[offset++];
+                  auto t32Val =
+                      rewriter.create<TritonGEN::FToTf32Op>(loc, f32Val)
+                          .getResult();
+                  matVal =
+                      tb.insert_element(dotOpTy, matVal, t32Val, tb.i32_val(k));
+
+                } else {
+                  matVal = tb.insert_element(dotOpTy, matVal, elems[offset++],
+                                             tb.i32_val(k));
+                }
               }
               vals[{b, i * repClusterOuter + repOuter,
                     j * repClusterInner + repInner}] = matVal;