[X86][AMX] support amxpreserve attribute in llvm.

[LuoYuanke]

Summary

Introduce 'amxpreserve' attribute to reduce save/restore AMX register
in register allocation.
Use this attribute to indicate that the specified function has no
caller-saved AMX registers. Compiler doesn't save and restore any
AMX register across function call. It is user's responsibility that
ensure there is no AMX register clobber in the function with "amxpreserve"
attribute.

Like 'no_caller_saved_registers', 'amxpreserve' attribute is not a
calling convention. In fact, it only overrides the decision of which
AMX registers should be saved by the caller.

For example:

.. code-block:: c

  __attribute__ ((amxpreserve ))
  void f () {
    ...
  }

  void bar () {
    ...
    f();
    ...
  }

In this case compiler doesn't save and restore AMX registers across the
call of f().

[jdoerfert]

Problems: no commit message, no RFC?, no lang ref update.

[LuoYuanke]

In D115202#3175738, @jdoerfert wrote:

Problems: no commit message, no RFC?, no lang ref update.

This is just for communication. I'll refine the patch and ask LLVM folks for review.

[jdoerfert]

In D115202#3175810, @LuoYuanke wrote:

In D115202#3175738, @jdoerfert wrote:

Problems: no commit message, no RFC?, no lang ref update.

This is just for communication. I'll refine the patch and ask LLVM folks for review.

That's fine. Please mark such patches as [WIP] or with a clear note in the commit message, or both.

[LuoYuanke]

That's fine. Please mark such patches as [WIP] or with a clear note in the commit message, or both.

Thanks for the suggestion. Done.

[LuoYuanke]

@jdoerfert, It is ready for review now.

[LuoYuanke]

update commit message.

[jdoerfert]

This needs a lang ref update and an RFC. If they exist, add links to them in the commit message.

llvm/include/llvm/AsmParser/LLToken.h

enum Kind {
   kw_uwtable,
   kw_vscale_range,
   kw_willreturn,
   kw_writeonly,
   kw_zeroext,
   kw_immarg,
   kw_byref,
   kw_mustprogress,
+  kw_amxpreserve,
 
   kw_type,
   kw_opaque,
 
   kw_comdat,
 
   // Comdat types
   kw_any,

llvm/include/llvm/Bitcode/LLVMBitCodes.h

enum AttributeKindCodes {
   ATTR_KIND_NO_CALLBACK = 71,
   ATTR_KIND_HOT = 72,
   ATTR_KIND_NO_PROFILE = 73,
   ATTR_KIND_VSCALE_RANGE = 74,
   ATTR_KIND_SWIFT_ASYNC = 75,
   ATTR_KIND_NO_SANITIZE_COVERAGE = 76,
   ATTR_KIND_ELEMENTTYPE = 77,
   ATTR_KIND_DISABLE_SANITIZER_INSTRUMENTATION = 78,
+  ATTR_KIND_AMXPRESERVE = 79,
 };
 
 enum ComdatSelectionKindCodes {
   COMDAT_SELECTION_KIND_ANY = 1,
   COMDAT_SELECTION_KIND_EXACT_MATCH = 2,
   COMDAT_SELECTION_KIND_LARGEST = 3,
   COMDAT_SELECTION_KIND_NO_DUPLICATES = 4,
   COMDAT_SELECTION_KIND_SAME_SIZE = 5,

llvm/include/llvm/IR/Attributes.td

 def WriteOnly : EnumAttr<"writeonly", [FnAttr, ParamAttr]>;
 
 /// Zero extended before/after call.
 def ZExt : EnumAttr<"zeroext", [ParamAttr, RetAttr]>;
 
 /// Function is required to make Forward Progress.
 def MustProgress : EnumAttr<"mustprogress", [FnAttr]>;
 
+/// Preserve AMX state across function call.
+def AMXPreserve: EnumAttr<"amxpreserve", [FnAttr]>;
+
 /// Target-independent string attributes.
 def LessPreciseFPMAD : StrBoolAttr<"less-precise-fpmad">;
 def NoInfsFPMath : StrBoolAttr<"no-infs-fp-math">;
 def NoNansFPMath : StrBoolAttr<"no-nans-fp-math">;
 def NoSignedZerosFPMath : StrBoolAttr<"no-signed-zeros-fp-math">;
 def UnsafeFPMath : StrBoolAttr<"unsafe-fp-math">;
 def NoJumpTables : StrBoolAttr<"no-jump-tables">;
 def NoInlineLineTables : StrBoolAttr<"no-inline-line-tables">;

llvm/lib/AsmParser/LLLexer.cpp

#define KEYWORD(STR)                                                           \
   KEYWORD(uwtable);
   KEYWORD(vscale_range);
   KEYWORD(willreturn);
   KEYWORD(writeonly);
   KEYWORD(zeroext);
   KEYWORD(immarg);
   KEYWORD(byref);
   KEYWORD(mustprogress);
+  KEYWORD(amxpreserve);
 
   KEYWORD(type);
   KEYWORD(opaque);
 
   KEYWORD(comdat);
 
   // Comdat types
   KEYWORD(any);

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

static Attribute::AttrKind getAttrFromCode(uint64_t Code) {
   case bitc::ATTR_KIND_UW_TABLE:
     return Attribute::UWTable;
   case bitc::ATTR_KIND_VSCALE_RANGE:
     return Attribute::VScaleRange;
   case bitc::ATTR_KIND_WILLRETURN:
     return Attribute::WillReturn;
   case bitc::ATTR_KIND_WRITEONLY:
     return Attribute::WriteOnly;
+  case bitc::ATTR_KIND_AMXPRESERVE:
+    return Attribute::AMXPreserve;
   case bitc::ATTR_KIND_Z_EXT:
     return Attribute::ZExt;
   case bitc::ATTR_KIND_IMMARG:
     return Attribute::ImmArg;
   case bitc::ATTR_KIND_SANITIZE_MEMTAG:
     return Attribute::SanitizeMemTag;
   case bitc::ATTR_KIND_PREALLOCATED:
     return Attribute::Preallocated;

llvm/lib/Bitcode/Writer/BitcodeWriter.cpp

static uint64_t getAttrKindEncoding(Attribute::AttrKind Kind) {
   case Attribute::SwiftAsync:
     return bitc::ATTR_KIND_SWIFT_ASYNC;
   case Attribute::UWTable:
     return bitc::ATTR_KIND_UW_TABLE;
   case Attribute::VScaleRange:
     return bitc::ATTR_KIND_VSCALE_RANGE;
   case Attribute::WillReturn:
     return bitc::ATTR_KIND_WILLRETURN;
+  case Attribute::AMXPreserve:
+    return bitc::ATTR_KIND_AMXPRESERVE;
   case Attribute::WriteOnly:
     return bitc::ATTR_KIND_WRITEONLY;
   case Attribute::ZExt:
     return bitc::ATTR_KIND_Z_EXT;
   case Attribute::ImmArg:
     return bitc::ATTR_KIND_IMMARG;
   case Attribute::SanitizeMemTag:
     return bitc::ATTR_KIND_SANITIZE_MEMTAG;

llvm/lib/Transforms/Utils/CodeExtractor.cpp

if (Attr.isStringAttribute()) {
       case Attribute::SExt:
       case Attribute::Speculatable:
       case Attribute::StackAlignment:
       case Attribute::StructRet:
       case Attribute::SwiftError:
       case Attribute::SwiftSelf:
       case Attribute::SwiftAsync:
       case Attribute::WillReturn:
+      case Attribute::AMXPreserve:
       case Attribute::WriteOnly:
       case Attribute::ZExt:
       case Attribute::ImmArg:
       case Attribute::ByRef:
       case Attribute::EndAttrKinds:
       case Attribute::EmptyKey:
       case Attribute::TombstoneKey:
         continue;

llvm/test/Bitcode/attributes.ll

 }
 
 ; CHECK: define void @f80() #50
 define void @f80() disable_sanitizer_instrumentation
 {
         ret void;
 }
 
+; CHECK: define void @f81() #51
+define void @f81() amxpreserve
+{
+  ret void
+}
+
 ; CHECK: attributes #0 = { noreturn }
 ; CHECK: attributes #1 = { nounwind }
 ; CHECK: attributes #2 = { readnone }
 ; CHECK: attributes #3 = { readonly }
 ; CHECK: attributes #4 = { noinline }
 ; CHECK: attributes #5 = { alwaysinline }
 ; CHECK: attributes #6 = { optsize }
 ; CHECK: attributes #7 = { ssp }

 ; CHECK: attributes #43 = { cold }
 ; CHECK: attributes #44 = { hot }
 ; CHECK: attributes #45 = { vscale_range(8,8) }
 ; CHECK: attributes #46 = { vscale_range(1,8) }
 ; CHECK: attributes #47 = { vscale_range(1,0) }
 ; CHECK: attributes #48 = { nosanitize_coverage }
 ; CHECK: attributes #49 = { noprofile }
 ; CHECK: attributes #50 = { disable_sanitizer_instrumentation }
+; CHECK: attributes #51 = { amxpreserve }
 ; CHECK: attributes #[[NOBUILTIN]] = { nobuiltin }