Added support for mma migration

Author: TejaX-Alaghari

clang/lib/DPCT/RulesAsm/AsmMigration.cpp

@@ -1301,6 +1301,83 @@ class SYCLGen : public SYCLGenBase {
     return SYCLGenSuccess();
   }
 
+  bool handle_mma(const InlineAsmInstruction *Inst) override {
+    if (Inst->getNumInputOperands() != 3)
+      return SYCLGenError();
+
+    if (!Inst->hasAttr(InstAttr::m16n8k16))
+      return SYCLGenError();
+
+    // Only row Layout is supported for of A matrix and
+    // only col Layout is supported for of B matrix
+    if (Inst->getAttr(3) != InstAttr::row ||
+        Inst->getAttr(4) != InstAttr::col) {
+      return SYCLGenError();
+    }
+
+    // Only f16 type is supported for A and B matrix data
+    const auto *AType = dyn_cast<InlineAsmBuiltinType>(Inst->getType(1));
+    const auto *BType = dyn_cast<InlineAsmBuiltinType>(Inst->getType(2));
+
+    std::string TypeStr;
+    if (!AType || !BType ||
+        (AType->getKind() != InlineAsmBuiltinType::f16 ||
+         BType->getKind() != InlineAsmBuiltinType::f16)) {
+      return SYCLGenError();
+    } else {
+      if (tryEmitType(TypeStr, AType))
+        return SYCLGenError();
+    }
+
+    const InlineAsmVectorExpr *VE =
+        dyn_cast<InlineAsmVectorExpr>(Inst->getOutputOperand());
+    if (VE && VE->getNumElements() != 4) {
+      return SYCLGenError();
+    }
+
+    OS() << MapNames::getDpctNamespace() << "experimental::matrix::mma";
+    OS() << "<" << TypeStr << ">(";
+
+    // Add D matrix address values to store the MAD result
+    for (unsigned Inst = 0; Inst != VE->getNumElements(); ++Inst) {
+      if (isa<InlineAsmDiscardExpr>(VE->getElement(Inst)))
+        continue;
+      OS() << "&";
+      if (emitStmt(VE->getElement(Inst)))
+        return SYCLGenError();
+      OS() << ", ";
+    }
+
+    // Add A, B & C matrix values to compute MAD
+    for (unsigned InputOp = 0; InputOp < Inst->getNumInputOperands();
+         InputOp++) {
+      if (VE = dyn_cast<InlineAsmVectorExpr>(Inst->getInputOperand(InputOp))) {
+        for (unsigned Inst = 0; Inst != VE->getNumElements(); ++Inst) {
+          if (isa<InlineAsmDiscardExpr>(VE->getElement(Inst)))
+            continue;
+          if (emitStmt(VE->getElement(Inst)))
+            return SYCLGenError();
+          OS() << ", ";
+        }
+      } else {
+        return SYCLGenError();
+      }
+    }
+
+    OS() << DpctGlobalInfo::getItem(GAS);
+    OS() << ");";
+
+    const auto *KernelDecl = getImmediateOuterFuncDecl(GAS);
+    if (KernelDecl) {
+      auto FuncInfo = DeviceFunctionDecl::LinkRedecls(KernelDecl);
+      if (FuncInfo)
+        FuncInfo->addSubGroupSizeRequest(32, GAS->getBeginLoc(),
+                                         DpctGlobalInfo::getSubGroup(GAS));
+    }
+
+    return SYCLGenSuccess();
+  }
+
   bool handle_prefetch(const InlineAsmInstruction *Inst) override {
     if (!DpctGlobalInfo::useExtPrefetch() || Inst->getNumInputOperands() != 1)
       return SYCLGenError();

clang/lib/DPCT/RulesAsm/Parser/AsmNodes.h

@@ -322,7 +322,7 @@ class InlineAsmInstruction : public InlineAsmStmt {
 
   /// This represents arrtibutes like: comparsion operator, rounding modifiers,
   /// ... e.g. instruction setp.eq.s32 has a comparsion operator 'eq'.
-  SmallSet<InstAttr, 4> Attributes;
+  SmallVector<InstAttr, 4> Attributes;
 
   /// This represents types in instruction, e.g. mov.u32.
   SmallVector<InlineAsmType *, 4> Types;
@@ -350,11 +350,11 @@ class InlineAsmInstruction : public InlineAsmStmt {
         OutputOp(Out), PredOutputOp(Pred), InputOps(InOps) {
     StateSpaces.insert(StateSpaces.begin(), AsmStateSpaces.begin(),
                        AsmStateSpaces.end());
-    Attributes.insert(Attrs.begin(), Attrs.end());
+    Attributes.insert(Attributes.begin(), Attrs.begin(), Attrs.end());
   }
 
   using attr_range =
-      llvm::iterator_range<SmallSet<InstAttr, 4>::const_iterator>;
+      llvm::iterator_range<SmallVector<InstAttr, 4>::const_iterator>;
   using type_range =
       llvm::iterator_range<SmallVector<InlineAsmType *, 4>::const_iterator>;
   using op_range =
@@ -369,12 +369,16 @@ class InlineAsmInstruction : public InlineAsmStmt {
   }
 
   template <typename... Ts> bool hasAttr(Ts... Attrs) const {
-    return (Attributes.contains(Attrs) || ...);
+    return (llvm::is_contained(Attributes, Attrs) || ...);
   }
   const InlineAsmIdentifierInfo *getOpcodeID() const { return Opcode; }
   asmtok::TokenKind getOpcode() const { return Opcode->getTokenID(); }
   ArrayRef<InlineAsmType *> getTypes() const { return Types; }
   const InlineAsmType *getType(unsigned I) const { return Types[I]; }
+  InstAttr getAttr(unsigned I) const {
+    assert(I < Attributes.size() && "Attributes index out of range");
+    return Attributes[I];
+  }
   unsigned getNumTypes() const { return Types.size(); }
   const InlineAsmExpr *getOutputOperand() const { return OutputOp; }
   const InlineAsmExpr *getPredOutputOperand() const { return PredOutputOp; }

clang/lib/DPCT/RulesAsm/Parser/AsmTokenKinds.def

@@ -274,6 +274,13 @@ MODIFIER(v2, ".v2")
 MODIFIER(v4, ".v4")
 MODIFIER(v8, ".v8")
 
+// Matrix modifiers
+MODIFIER(row, ".row")
+MODIFIER(col, ".col")
+
+// Matrix shape
+MODIFIER(m16n8k16, ".m16n8k16")
+
 STATE_SPACE(reg,    ".reg")
 STATE_SPACE(sreg,   ".sreg")
 STATE_SPACE(const,  ".const")
@@ -418,6 +425,7 @@ MODIFIER(rc8, ".rc8")
 MODIFIER(ecl, ".ecl")
 MODIFIER(ecr, ".ecr")
 MODIFIER(rc16, ".rc16")
+MODIFIER(aligned, ".aligned")
 
 #undef LINKAGE
 #undef TARGET

clang/runtime/dpct-rt/include/dpct/math.hpp

@@ -2055,6 +2055,80 @@ class joint_matrix {
   matrix_accessor x;
   const size_t num_elements;
 };
+
+/// Multiplies 2 16x16 & 16x8 matrices and accumulates the result to a 16x8 b32
+/// matrix
+/// \tparam [in] MulType The type of the multiplication result
+/// \tparam [in] ABType The type of the input matrices
+/// \tparam [in] CDType The type of the output matrix
+/// \param [in] aTrans Indicates whether the 1st matrix to be transposed
+/// \param [in] bTrans Indicates whether the 2nd matrix to be transposed
+/// \param [in] d0 The 1st element to be written to the output D matrix
+/// \param [in] d1 The 2nd element to be written to the output D matrix
+/// \param [in] d2 The 3rd element to be written to the output D matrix
+/// \param [in] d3 The 4th element to be written to the output D matrix
+/// \param [in] a0 The 1st element from A matrix to be multiplied with B matrix
+/// \param [in] a1 The 2nd element from A matrix to be multiplied with B matrix
+/// \param [in] a2 The 3rd element from A matrix to be multiplied with B matrix
+/// \param [in] a3 The 4th element from A matrix to be multiplied with B matrix
+/// \param [in] b0 The 1st element from B matrix to be multiplied with A matrix
+/// \param [in] b1 The 2nd element from B matrix to be multiplied with A matrix
+/// \param [in] c0 The 1st element from C matrix to be added with d0
+/// \param [in] c1 The 2nd element from C matrix to be added with d1
+/// \param [in] c2 The 3rd element from C matrix to be added with d2
+/// \param [in] c3 The 4th element from C matrix to be added with d3
+/// \param [in] item The sycl::nd_item index space class
+template <typename MulType, typename ABType, typename CDType, typename ItemT>
+__attribute__((optnone)) void
+mma(CDType *d0, CDType *d1, CDType *d2, CDType *d3, ABType a0, ABType a1,
+    ABType a2, ABType a3, ABType b0, ABType b1, CDType c0, CDType c1, CDType c2,
+    CDType c3, const ItemT &item) {
+  int lane = item.get_sub_group().get_local_linear_id();
+
+  short ROW_LOAD_OFFSET = 4 * (lane / 4);
+  short COL_LOAD_OFFSET = 8 * (lane % 4);
+
+  ABType recv_a[4 * 4], recv_b[4 * 4];
+  for (int i = 0; i < 4; i++) {
+    recv_a[0 * 4 + i] = dpct::select_from_sub_group(item.get_sub_group(), a0,
+                                                    ROW_LOAD_OFFSET + i);
+    recv_a[1 * 4 + i] = dpct::select_from_sub_group(item.get_sub_group(), a2,
+                                                    ROW_LOAD_OFFSET + i);
+    recv_a[2 * 4 + i] = dpct::select_from_sub_group(item.get_sub_group(), a1,
+                                                    ROW_LOAD_OFFSET + i);
+    recv_a[3 * 4 + i] = dpct::select_from_sub_group(item.get_sub_group(), a3,
+                                                    ROW_LOAD_OFFSET + i);
+
+    recv_b[0 * 4 + i] = dpct::select_from_sub_group(item.get_sub_group(), b0,
+                                                    COL_LOAD_OFFSET + i);
+    recv_b[1 * 4 + i] = dpct::select_from_sub_group(item.get_sub_group(), b1,
+                                                    COL_LOAD_OFFSET + i);
+    recv_b[2 * 4 + i] = dpct::select_from_sub_group(item.get_sub_group(), b0,
+                                                    COL_LOAD_OFFSET + 4 + i);
+    recv_b[3 * 4 + i] = dpct::select_from_sub_group(item.get_sub_group(), b1,
+                                                    COL_LOAD_OFFSET + 4 + i);
+  }
+
+  auto *a = reinterpret_cast<MulType *>(recv_a);
+  auto *b = reinterpret_cast<MulType *>(recv_b);
+  for (int i = 0; i < 16; i++) {
+    auto a0 = static_cast<CDType>(a[i]);
+    auto a1 = static_cast<CDType>(a[i + 16]);
+    auto b0 = static_cast<CDType>(b[i]);
+    auto b1 = static_cast<CDType>(b[i + 16]);
+
+    c0 += a0 * b0;
+    c1 += a0 * b1;
+    c2 += a1 * b0;
+    c3 += a1 * b1;
+  }
+
+  *d0 = c0;
+  *d1 = c1;
+  *d2 = c2;
+  *d3 = c3;
+}
+
 } // namespace matrix
 } // namespace experimental
 

clang/test/dpct/asm/mma.cu

@@ -0,0 +1,45 @@
+// UNSUPPORTED: cuda-8.0, cuda-9.0, cuda-9.1, cuda-9.2, cuda-10.0, cuda-10.1, cuda-10.2
+// UNSUPPORTED: v8.0, v9.0, v9.1, v9.2, v10.0, v10.1, v10.2
+// RUN: dpct --format-range=none -out-root %T/mma %s --cuda-include-path="%cuda-path/include" -- -std=c++14 -x cuda --cuda-host-only
+// RUN: FileCheck %s --match-full-lines --input-file %T/mma/mma.dp.cpp
+// RUN: %if build_lit %{icpx -c -DNO_BUILD_TEST -fsycl %T/mma/mma.dp.cpp -o %T/mma/mma.dp.o %}
+
+// clang-format off
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+
+/*
+mma.sync.aligned.m16n8k16.alayout.blayout.dtype.f16.f16.ctype d, a, b, c;
+
+Below are the currenly supported configurations:
+
+.alayout = {.row};
+.blayout = {.col};
+.ctype   = {.f32};
+.dtype   = {.f32};
+*/
+
+__global__ void mma_kernel() {
+  int a[4];
+  int b[2];
+  float c[4];
+
+  // CHECK: dpct::experimental::matrix::mma<sycl::half>(&c[0], &c[1], &c[2], &c[3], a[0], a[1], a[2], a[3], b[0], b[1], c[0], c[1], c[2], c[3], item_ct1);
+  asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+        " { %0, %1, %2, %3 }, "
+        " { %4, %5, %6, %7 }, "
+        " { %8, %9 }, "
+        " { %0, %1, %2, %3 };"
+        : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
+        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]),
+          "r"(b[0]), "r"(b[1]));
+}
+
+
+int main () {
+  // CHECK: [=](sycl::nd_item<3> item_ct1) {{\[\[}}sycl::reqd_sub_group_size(32){{\]\]}} {
+  mma_kernel<<<1, 32>>>();
+
+  return 0;
+}
+// clang-format on

docs/dev_guide/api-mapping-status/ASM_API_migration_status.csv

@@ -54,7 +54,7 @@ max,YES,
 mbarrier,NO,
 membar,YES, Partial
 min,YES,
-mma,NO,
+mma,YES, Partial
 mov,YES,
 movmatrix,NO,
 mul,YES, Partial