Add xevm integration tests

Author: akroviakov

include/gc/Dialect/LLVMIR/XeVMOps.td

@@ -221,6 +221,76 @@ def XeVM_BlockPrefetch2dOp : XeVM_Op<"blockprefetch2d">,
   let hasVerifier = 1;
 }
 
+def XeVM_MatrixElemType : AnyTypeOf<[AnyI8, AnyI16, AnyI32, F32, F16, BF16]>;
+
+/// Enum attribute of the different precision types.
+def XeVM_PrecisionTypeAttr : I32EnumAttr<"PrecisionType",
+  "XeVM precision type",
+  [
+    I32EnumAttrCase<"UNUSED", 0,  "unused">,
+    I32EnumAttrCase<"U8",     1,  "u8">,
+    I32EnumAttrCase<"U4",     2,  "u4">,
+    I32EnumAttrCase<"U2",     3,  "u2">,
+    I32EnumAttrCase<"S8",     4,  "i8">,
+    I32EnumAttrCase<"S4",     5,  "i4">,
+    I32EnumAttrCase<"S2",     6,  "i2">,
+    I32EnumAttrCase<"BF8",    7,  "bf8">,
+    I32EnumAttrCase<"TF32",   8,  "tf32">,
+    I32EnumAttrCase<"BF16",   9,  "bf16">,
+    I32EnumAttrCase<"FP16",   10, "f16">
+  ]> {
+  let cppNamespace = "::mlir::xevm";
+}
+
+def XeVM_DPASOp : XeVM_Op<"dpas">,
+  Results<(outs FixedVectorOf<[XeVM_MatrixElemType]>:$d)>,
+  Arguments<(ins
+    FixedVectorOfRankAndType<[1], [XeVM_MatrixElemType]>:$c,
+    FixedVectorOfRankAndType<[1], [XeVM_MatrixElemType]>:$a,
+    FixedVectorOfRankAndType<[1], [XeVM_MatrixElemType]>:$b,
+    XeVM_PrecisionTypeAttr:$pa,
+    XeVM_PrecisionTypeAttr:$pb,
+    I32Attr:$rc
+  )> {
+
+  let summary = "Matrix multiply-add";
+
+  let description = [{
+    The `xevm.dpas` operation is a matrix multiplication plus accumulation:
+
+      D = C + A x B
+
+      where the A, B, C input matrices and the result D have shapes:
+        D : MxN
+        C : MxN
+        A : MxK
+        B : KxN
+        
+        Shape restrictions:
+        M : must be 1, 2, 4, or 8
+        N : fixed execution size, must be 16
+        K : systolic_depth * OPS_PER_CHAN
+            OPS_PER_CHAN
+              1 : for TF32
+              2 : for 16-bit precision(BF, HF)
+              4 : for 8-bit precision (FP8, UB, B)
+              8 : for less-then 8 bit precision (U4/S4, U2/S2).
+
+            If systolic_depth is 8, K would be 8, 16, 32, or 64 (based on OPS_PER_CHAN).
+    $a, $b, $c, $d - matrix A, B, C, D, respectively
+    $pa, $pb - precision of matrix A and B resepectively
+    $rc - repeat count
+
+    Further restrictions as well as more details can be found here:
+    https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_matrix_multiply_accumulate.html
+  }];
+
+  let assemblyFormat = [{
+    operands ` ` `{` `pa` `=` $pa `,` `pb` `=` $pb `,` `rc` `=` $rc `}` attr-dict `:` functional-type(operands, results)
+  }];
+
+  // let hasVerifier = 1;
+}
 
 def XeVM_TargetAttr : XeVM_Attr<"XeVMTarget", "target"> {
   let description = [{

include/gc/ExecutionEngine/Driver/Driver.h

@@ -18,7 +18,7 @@ namespace mlir {
 class DialectRegistry;
 namespace gc {
 
-const DialectRegistry &initCompilerAndGetDialects();
+DialectRegistry &initCompilerAndGetDialects();
 
 // the pointers to XXXMemRefType
 using GeneralMemrefPtr = void *;

lib/gc/Conversion/XeVMToLLVM/XeVMToLLVM.cpp

@@ -11,6 +11,7 @@
 #include "gc/Dialect/LLVMIR/XeVMDialect.h"
 #include "mlir/Conversion/ConvertToLLVM/ToLLVMInterface.h"
 #include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Pass/Pass.h"
@@ -53,6 +54,8 @@ static constexpr LLVMFuncAttributeOptions noUnwindAttrs = {
     false, true, false, {}};
 static constexpr LLVMFuncAttributeOptions noUnwindWillReturnAttrs = {
     false, true, true, {}};
+static constexpr LLVMFuncAttributeOptions convergentNoUnwindWillReturnAttrs = {
+    true, true, true, {}};
 
 std::string getTypeMangling(Type ty, bool isUnsigned = false) {
   return TypeSwitch<Type, std::string>(ty)
@@ -79,6 +82,31 @@ std::string getTypeMangling(Type ty, bool isUnsigned = false) {
       });
 }
 
+std::string mangle(StringRef baseName, ArrayRef<Type> types,
+                   ArrayRef<bool> isUnsigned = {}) {
+  assert((isUnsigned.empty() || isUnsigned.size() == types.size()) &&
+         "Signedness info doesn't match");
+  std::string s;
+  llvm::raw_string_ostream os(s);
+  llvm::SmallDenseMap<Type, unsigned> substitutions;
+  os << "_Z" << baseName.size() << baseName;
+  for (auto [idx, type] : llvm::enumerate(types)) {
+    auto it = substitutions.find(type);
+    if (it != substitutions.end()) {
+      os << "S";
+      // First substitution is `S_`, second is `S0_`, and so on.
+      if (unsigned firstIdx = it->getSecond(); firstIdx > 0)
+        os << firstIdx - 1;
+      os << "_";
+    } else {
+      if (!type.isIntOrFloat())
+        substitutions[type] = substitutions.size();
+      os << getTypeMangling(type, isUnsigned.empty() ? false : isUnsigned[idx]);
+    }
+  }
+  return os.str();
+}
+
 template <typename OpType>
 static std::optional<ArrayAttr>
 getCacheControlMetadata(ConversionPatternRewriter &rewriter, OpType op,
@@ -115,13 +143,15 @@ getCacheControlMetadata(ConversionPatternRewriter &rewriter, OpType op,
   return rewriter.getArrayAttr(combinedAttrs);
 }
 
-static LLVM::CallOp
-createDeviceFunctionCall(ConversionPatternRewriter &rewriter,
-                         StringRef funcName, Type retType,
-                         ArrayRef<Type> argTypes, ArrayRef<Value> args,
-                         ArrayRef<std::pair<unsigned, StringRef>> paramAttrs,
-                         LLVMFuncAttributeOptions funcAttributeOptions) {
-  auto moduleOp = rewriter.getBlock()->getParent()->getParentOfType<ModuleOp>();
+static LLVM::CallOp createDeviceFunctionCall(
+    ConversionPatternRewriter &rewriter, StringRef funcName, Type retType,
+    ArrayRef<Type> argTypes, ArrayRef<Value> args,
+    mlir::ArrayRef<std::pair<unsigned, mlir::StringRef>> paramAttrs,
+    LLVMFuncAttributeOptions funcAttributeOptions) {
+  auto moduleOp = rewriter.getBlock()
+                      ->getParentOp()
+                      ->getParentWithTrait<OpTrait::SymbolTable>();
+  assert(moduleOp && "Expecting module");
   MLIRContext *ctx = rewriter.getContext();
   Location loc = UnknownLoc::get(ctx);
 
@@ -144,6 +174,96 @@ createDeviceFunctionCall(ConversionPatternRewriter &rewriter,
   return callOp;
 }
 
+class DPASToOCLPattern : public OpConversionPattern<xevm::DPASOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(xevm::DPASOp op, xevm::DPASOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    constexpr uint32_t bitWidthPackedA{16};
+    constexpr uint32_t bitWidthPackedB{32};
+    auto loc = op.getLoc();
+
+    auto castIfNeeded = [&](Value val, Type packedType) -> Value {
+      VectorType origTy = cast<VectorType>(val.getType());
+      const uint32_t vecBitSize =
+          origTy.getNumElements() *
+          origTy.getElementType().getIntOrFloatBitWidth();
+      VectorType newTy = VectorType::get(
+          vecBitSize / packedType.getIntOrFloatBitWidth(), packedType);
+      if (origTy != newTy)
+        val = rewriter.create<LLVM::BitcastOp>(loc, newTy, val);
+      return val;
+    };
+
+    Value a = op.getA();
+    Type packedAType = (op.getPa() == xevm::PrecisionType::TF32)
+                           ? cast<Type>(rewriter.getF32Type())
+                           : rewriter.getIntegerType(bitWidthPackedA);
+    a = castIfNeeded(a, packedAType);
+
+    Value b = op.getB();
+    Type packedBType = (op.getPb() == xevm::PrecisionType::TF32)
+                           ? cast<Type>(rewriter.getF32Type())
+                           : rewriter.getIntegerType(bitWidthPackedB);
+    b = castIfNeeded(b, packedBType);
+
+    Value c = op.getC();
+    VectorType cOrigTy = cast<VectorType>(c.getType());
+    assert(cOrigTy == op->getResultTypes()[0] &&
+           "Accumulator and result type mismatch");
+    // OCL builtins encode bfloat16 as int16
+    VectorType cTy =
+        cOrigTy.getElementType().isBF16()
+            ? VectorType::get(cOrigTy.getShape(), rewriter.getIntegerType(16))
+            : cOrigTy;
+    if (cOrigTy != cTy)
+      c = rewriter.create<LLVM::BitcastOp>(loc, cTy, c);
+
+    constexpr int32_t systolicDepth{8};
+    std::string fnName =
+        llvm::formatv("intel_sub_group_{0}_{1}_matrix_mad_k{2}",
+                      stringifyPrecisionType(op.getPa()).str(),
+                      stringifyPrecisionType(op.getPb()).str(),
+                      systolicDepth * getNumOperandsPerDword(op.getPa()))
+            .str();
+    SmallVector<Type> argTypes{a.getType(), b.getType(), cTy};
+    fnName = mangle(fnName, argTypes);
+    SmallVector<Value> args{a, b, c};
+
+    auto memAttr = rewriter.getAttr<LLVM::MemoryEffectsAttr>(
+        /*other=*/LLVM::ModRefInfo::NoModRef,
+        /*argMem=*/LLVM::ModRefInfo::NoModRef,
+        /*inaccessibleMem=*/LLVM::ModRefInfo::NoModRef);
+    auto funcAttrs = convergentNoUnwindWillReturnAttrs;
+    funcAttrs.memEffectsAttr = memAttr;
+    Value result = createDeviceFunctionCall(rewriter, fnName, cTy, argTypes,
+                                            args, {}, funcAttrs)
+                       ->getResult(0);
+
+    if (cOrigTy != cTy)
+      result = rewriter.create<LLVM::BitcastOp>(loc, cOrigTy, result);
+
+    rewriter.replaceOp(op, result);
+    return success();
+  }
+
+private:
+  static unsigned getNumOperandsPerDword(xevm::PrecisionType pTy) {
+    switch (pTy) {
+    case xevm::PrecisionType::TF32:
+      return 1;
+    case xevm::PrecisionType::BF16:
+    case xevm::PrecisionType::FP16:
+      return 2;
+    case xevm::PrecisionType::U8:
+    case xevm::PrecisionType::S8:
+      return 4;
+    default:
+      llvm_unreachable("unsupported xevm::PrecisionType");
+    }
+  }
+};
+
 template <typename OpType>
 class LoadStorePrefetchToOCLPattern : public OpConversionPattern<OpType> {
   using OpConversionPattern<OpType>::OpConversionPattern;
@@ -291,10 +411,11 @@ struct ConvertXeVMToLLVMPass
 //===----------------------------------------------------------------------===//
 
 void mlir::populateXeVMToLLVMConversionPatterns(RewritePatternSet &patterns) {
-  patterns.add<LoadStorePrefetchToOCLPattern<BlockLoad2dOp>,
-               LoadStorePrefetchToOCLPattern<BlockStore2dOp>,
-               LoadStorePrefetchToOCLPattern<BlockPrefetch2dOp>>(
-      patterns.getContext());
+  patterns
+      .add<LoadStorePrefetchToOCLPattern<BlockLoad2dOp>,
+           LoadStorePrefetchToOCLPattern<BlockStore2dOp>,
+           LoadStorePrefetchToOCLPattern<BlockPrefetch2dOp>, DPASToOCLPattern>(
+          patterns.getContext());
 }
 
 //===----------------------------------------------------------------------===//

lib/gc/ExecutionEngine/Driver/CMakeLists.txt

@@ -25,6 +25,7 @@ else()
     MLIRToLLVMIRTranslationRegistration
   )
 endif()
+get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 
 set(GC_PASSES GcInterface GcPasses)
 if(GC_ENABLE_IMEX)
@@ -38,6 +39,8 @@ gc_add_mlir_library(GcJitWrapper
     ${MLIR_LINK_COMPONENTS}
     ${dialect_libs}
     ${conversion_libs}
+    ${extension_libs}
     ${GC_PASSES}
     GcAnalysis
+    MLIRXeVMToLLVMIRTranslation
   )

lib/gc/ExecutionEngine/Driver/Driver.cpp

@@ -11,8 +11,12 @@
 #ifdef GC_HAS_ONEDNN_DIALECT
 #include "gc/Dialect/OneDNNGraph/OneDNNGraphDialect.h"
 #endif
+#include "gc/Conversion/Passes.h"
+#include "gc/Target/LLVM/XeVM/Target.h"
+#include "gc/Target/LLVMIR/Dialect/XeVM/XeVMToLLVMIRTranslation.h"
 #include "gc/Transforms/Passes.h"
 #include "mlir/InitAllDialects.h"
+#include "mlir/InitAllExtensions.h"
 #include "mlir/InitAllPasses.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Target/LLVMIR/Dialect/All.h"
@@ -26,22 +30,29 @@ namespace gc {
 static DialectRegistry initDialects() {
   mlir::registerAllPasses();
   mlir::gc::registerGraphCompilerPasses();
+  mlir::registerGCConversionPasses();
   mlir::cpuruntime::registerCPURuntimePasses();
   mlir::DialectRegistry registry;
   registry.insert<mlir::cpuruntime::CPURuntimeDialect>();
   mlir::registerAllDialects(registry);
   mlir::cpuruntime::registerConvertCPURuntimeToLLVMInterface(registry);
+  mlir::registerAllExtensions(registry);
+  // Adds missing `LLVMTranslationDialectInterface` registration for dialect for
+  // gpu.module op
+  mlir::registerAllToLLVMIRTranslations(registry);
+  mlir::registerConvertXeVMToLLVMInterface(registry);
+  mlir::registerXeVMDialectTranslation(registry);
+  mlir::xevm::registerXeVMTargetInterfaceExternalModels(registry);
 #ifdef GC_HAS_ONEDNN_DIALECT
   registry.insert<mlir::onednn_graph::OneDNNGraphDialect>();
 #endif
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
   llvm::InitializeNativeTargetAsmParser();
-  mlir::registerAllToLLVMIRTranslations(registry);
   return registry;
 }
 
-const DialectRegistry &initCompilerAndGetDialects() {
+DialectRegistry &initCompilerAndGetDialects() {
   static DialectRegistry reg = initDialects();
   return reg;
 }