!200 [SME][matrix_type] lower matrix_type with ARM SME/SVE instructions

From: @chenzheng1030 
Reviewed-by: @eastB233 
Signed-off-by: @eastB233

Author: openeuler-ci-bot

clang/include/clang/AST/Type.h

@@ -4176,6 +4176,21 @@ class MatrixType : public Type, public llvm::FoldingSetNode {
            (T->isRealType() && !T->isBooleanType() && !T->isEnumeralType());
   }
 
+  static bool isValidTypeForSME(QualType T) {
+    if (!isValidElementType(T))
+      return false;
+
+    if (!isa<BuiltinType>(T))
+      return false;
+
+    // AArch64 can not do vector operations like fma/add/sub for __bf16.
+    if (T->isBFloat16Type())
+      return false;
+
+    return cast<BuiltinType>(T)->isFloatingPoint() ||
+           cast<BuiltinType>(T)->isInteger();
+  }
+
   bool isSugared() const { return false; }
   QualType desugar() const { return QualType(this, 0); }
 

clang/lib/CodeGen/CGBuiltin.cpp

@@ -3976,12 +3976,18 @@ RValue CodeGenFunction::EmitBuiltinExpr(const GlobalDecl GD, unsigned BuiltinID,
         *this, E, llvm::Intrinsic::vector_reduce_and, "rdx.and"));
 
   case Builtin::BI__builtin_matrix_transpose: {
-    auto *MatrixTy = E->getArg(0)->getType()->castAs<ConstantMatrixType>();
-    Value *MatValue = EmitScalarExpr(E->getArg(0));
+    auto *MatrixValue = E->getArg(0);
+    auto *MatrixTy = MatrixValue->getType()->castAs<ConstantMatrixType>();
+    Value *MatValue = EmitScalarExpr(MatrixValue);
     MatrixBuilder MB(Builder);
-    Value *Result = MB.CreateMatrixTranspose(MatValue, MatrixTy->getNumRows(),
-                                             MatrixTy->getNumColumns());
-    return RValue::get(Result);
+
+    if (!getContext().getTargetInfo().hasFeature("sme") ||
+        !MatrixType::isValidTypeForSME(MatrixTy->getElementType()))
+      return RValue::get(MB.CreateMatrixTranspose(
+          MatValue, MatrixTy->getNumRows(), MatrixTy->getNumColumns()));
+
+    return RValue::get(MB.CreateSMEMatrixTranspose(
+        MatValue,MatrixTy->getNumRows(),MatrixTy->getNumColumns()));
   }
 
   case Builtin::BI__builtin_matrix_column_major_load: {

clang/lib/CodeGen/CGExprScalar.cpp

@@ -764,10 +764,23 @@ class ScalarExprEmitter
       auto *RHSMatTy = dyn_cast<ConstantMatrixType>(
           BO->getRHS()->getType().getCanonicalType());
       CodeGenFunction::CGFPOptionsRAII FPOptsRAII(CGF, Ops.FPFeatures);
-      if (LHSMatTy && RHSMatTy)
-        return MB.CreateMatrixMultiply(Ops.LHS, Ops.RHS, LHSMatTy->getNumRows(),
-                                       LHSMatTy->getNumColumns(),
-                                       RHSMatTy->getNumColumns());
+      if (LHSMatTy && RHSMatTy) {
+        // Note that SME only has non-widening MOPA for float32 and float64, so
+        // only these two types have native SME matmul operations. For other
+        // types, SVE version is used. We hope that SVE version is better than
+        // default NEON or scalar version.
+        auto Ty = LHSMatTy->getElementType();
+        if (!CGF.getContext().getTargetInfo().hasFeature("sme") ||
+            !MatrixType::isValidTypeForSME(Ty))
+          return MB.CreateMatrixMultiply(
+              Ops.LHS, Ops.RHS, LHSMatTy->getNumRows(),
+              LHSMatTy->getNumColumns(), RHSMatTy->getNumColumns());
+        assert(isa<BuiltinType>(Ty) && "SME types should be BuiltinType.");
+        return MB.CreateSMEMatrixMultiply(
+            Ops.LHS, Ops.RHS, LHSMatTy->getNumRows(), LHSMatTy->getNumColumns(),
+            RHSMatTy->getNumColumns(),
+            cast<BuiltinType>(Ty)->isSignedInteger());
+      }
       return MB.CreateScalarMultiply(Ops.LHS, Ops.RHS);
     }
 
@@ -4170,7 +4183,16 @@ Value *ScalarExprEmitter::EmitAdd(const BinOpInfo &op) {
   if (op.Ty->isConstantMatrixType()) {
     llvm::MatrixBuilder MB(Builder);
     CodeGenFunction::CGFPOptionsRAII FPOptsRAII(CGF, op.FPFeatures);
-    return MB.CreateAdd(op.LHS, op.RHS);
+
+    auto *MatTy = cast<ConstantMatrixType>(op.E->getType().getCanonicalType());
+    auto Ty = MatTy->getElementType();
+    if (!CGF.getContext().getTargetInfo().hasFeature("sme") ||
+        !MatrixType::isValidTypeForSME(Ty))
+      return MB.CreateAdd(op.LHS, op.RHS);
+    assert(isa<BuiltinType>(Ty) && "SME types should be BuiltinType.");
+    return MB.CreateSMEMatrixBinOp(
+        op.LHS, op.RHS, MatTy->getNumRows(), MatTy->getNumColumns(),
+        cast<BuiltinType>(Ty)->isSignedInteger(), "add");
   }
 
   if (op.Ty->isUnsignedIntegerType() &&
@@ -4326,7 +4348,16 @@ Value *ScalarExprEmitter::EmitSub(const BinOpInfo &op) {
     if (op.Ty->isConstantMatrixType()) {
       llvm::MatrixBuilder MB(Builder);
       CodeGenFunction::CGFPOptionsRAII FPOptsRAII(CGF, op.FPFeatures);
-      return MB.CreateSub(op.LHS, op.RHS);
+      auto *MatTy =
+          cast<ConstantMatrixType>(op.E->getType().getCanonicalType());
+      auto Ty = MatTy->getElementType();
+      if (!CGF.getContext().getTargetInfo().hasFeature("sme") ||
+          !MatrixType::isValidTypeForSME(Ty))
+        return MB.CreateSub(op.LHS, op.RHS);
+      assert(isa<BuiltinType>(Ty) && "SME types should be BuiltinType.");
+      return MB.CreateSMEMatrixBinOp(
+          op.LHS, op.RHS, MatTy->getNumRows(), MatTy->getNumColumns(),
+          cast<BuiltinType>(Ty)->isSignedInteger(), "sub");
     }
 
     if (op.Ty->isUnsignedIntegerType() &&

clang/lib/CodeGen/Targets/AArch64.cpp

@@ -294,6 +294,10 @@ AArch64ABIInfo::classifyArgumentType(QualType Ty, bool IsVariadic,
   if (isIllegalVectorType(Ty))
     return coerceIllegalVector(Ty);
 
+  // Always pass the matrix type via memory.
+  if (Ty->isMatrixType())
+    return getNaturalAlignIndirect(Ty, false);
+
   if (!isAggregateTypeForABI(Ty)) {
     // Treat an enum type as its underlying type.
     if (const EnumType *EnumTy = Ty->getAs<EnumType>())
@@ -393,6 +397,10 @@ ABIArgInfo AArch64ABIInfo::classifyReturnType(QualType RetTy,
       return coerceIllegalVector(RetTy);
   }
 
+  // Always return the matrix type via memory.
+  if (RetTy->isMatrixType())
+    return getNaturalAlignIndirect(RetTy);
+
   // Large vector types should be returned via memory.
   if (RetTy->isVectorType() && getContext().getTypeSize(RetTy) > 128)
     return getNaturalAlignIndirect(RetTy);

clang/lib/Driver/ToolChain.cpp

@@ -1112,6 +1112,10 @@ ToolChain::RuntimeLibType ToolChain::GetRuntimeLibType(
     runtimeLibType = GetDefaultRuntimeLibType();
   }
 
+  const llvm::Triple::ArchType Arch = getArch();
+  if (Arch == llvm::Triple::aarch64 && Args.hasArg(options::OPT_fenable_matrix))
+    runtimeLibType = ToolChain::RLT_CompilerRT;
+
   return *runtimeLibType;
 }