[Clang][Attr] Introduce the `assume` function attribute

The `assume` attribute is a way to provide additional, arbitrary
information to the optimizer. For now, assumptions are restricted to
strings which will be accumulated for a function and emitted as comma
separated string function attribute. The key of the LLVM-IR function
attribute is `llvm.assume`. Similar to `llvm.assume` and
`__builtin_assume`, the `assume` attribute provides a user defined
assumption to the compiler.

A follow up patch will introduce an LLVM-core API to query the
assumptions attached to a function. We also expect to add more options,
e.g., expression arguments, to the `assume` attribute later on.

The `omp [begin] asssumes` pragma will leverage this attribute and
expose the functionality in the absence of OpenMP.

Reviewed By: aaron.ballman

Differential Revision: https://reviews.llvm.org/D91979

clang/docs/LanguageExtensions.rst

@@ -1740,6 +1740,8 @@ implemented directly in terms of :ref:`extended vector support
 <langext-vectors>` instead of builtins, in order to reduce the number of
 builtins that we need to implement.
 
+.. _langext-__builtin_assume:
+
 ``__builtin_assume``
 ------------------------------
 

clang/include/clang/Basic/Attr.td

@@ -3527,6 +3527,14 @@ def OMPDeclareVariant : InheritableAttr {
   }];
 }
 
+def Assumption : InheritableAttr {
+  let Spellings = [Clang<"assume">];
+  let Subjects = SubjectList<[Function, ObjCMethod]>;
+  let InheritEvenIfAlreadyPresent = 1;
+  let Documentation = [AssumptionDocs];
+  let Args = [StringArgument<"Assumption">];
+}
+
 def InternalLinkage : InheritableAttr {
   let Spellings = [Clang<"internal_linkage">];
   let Subjects = SubjectList<[Var, Function, CXXRecord]>;

clang/include/clang/Basic/AttrDocs.td

@@ -3992,6 +3992,41 @@ For more information see
 }];
 }
 
+def AssumptionDocs : Documentation {
+  let Category = DocCatFunction;
+  let Heading = "assume";
+  let Content = [{
+Clang supports the ``__attribute__((assume("assumption")))`` attribute to
+provide additional information to the optimizer. The string-literal, here
+"assumption", will be attached to the function declaration such that later
+analysis and optimization passes can assume the "assumption" to hold.
+This is similar to :ref:`__builtin_assume <langext-__builtin_assume>` but
+instead of an expression that can be assumed to be non-zero, the assumption is
+expressed as a string and it holds for the entire function.
+
+A function can have multiple assume attributes and they propagate from prior
+declarations to later definitions. Multiple assumptions are aggregated into a
+single comma separated string. Thus, one can provide multiple assumptions via
+a comma separated string, i.a.,
+``__attribute__((assume("assumption1,assumption2")))``.
+
+While LLVM plugins might provide more assumption strings, the default LLVM
+optimization passes are aware of the following assumptions:
+
+  .. code-block:: none
+
+    "omp_no_openmp"
+    "omp_no_openmp_routines"
+    "omp_no_parallelism"
+
+The OpenMP standard defines the meaning of OpenMP assumptions ("omp_XYZ" is
+spelled "XYZ" in the `OpenMP 5.1 Standard`_).
+
+.. _`OpenMP 5.1 Standard`: https://www.openmp.org/spec-html/5.1/openmpsu37.html#x56-560002.5.2
+
+}];
+}
+
 def NoStackProtectorDocs : Documentation {
   let Category = DocCatFunction;
   let Content = [{

clang/include/clang/Basic/DiagnosticGroups.td

@@ -19,6 +19,8 @@ def Implicit : DiagGroup<"implicit", [
 def ODR : DiagGroup<"odr">;
 def : DiagGroup<"abi">;
 def AbsoluteValue : DiagGroup<"absolute-value">;
+def MisspelledAssumption : DiagGroup<"misspelled-assumption">;
+def UnknownAssumption : DiagGroup<"unknown-assumption">;
 def AddressOfTemporary : DiagGroup<"address-of-temporary">;
 def : DiagGroup<"aggregate-return">;
 def GNUAlignofExpression : DiagGroup<"gnu-alignof-expression">;

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -740,6 +740,13 @@ def note_strncat_wrong_size : Note<
 def warn_assume_side_effects : Warning<
   "the argument to %0 has side effects that will be discarded">,
   InGroup<DiagGroup<"assume">>;
+def warn_assume_attribute_string_unknown : Warning<
+  "unknown assumption string '%0'; attribute is potentially ignored">,
+  InGroup<UnknownAssumption>;
+def warn_assume_attribute_string_unknown_suggested : Warning<
+  "unknown assumption string '%0' may be misspelled; attribute is potentially "
+  "ignored, did you mean '%1'?">,
+  InGroup<MisspelledAssumption>;
 
 def warn_builtin_chk_overflow : Warning<
   "'%0' will always overflow; destination buffer has size %1,"

clang/include/clang/Sema/Sema.h

@@ -58,8 +58,8 @@
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallBitVector.h"
-#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"

clang/lib/CodeGen/CGCall.cpp

@@ -31,6 +31,7 @@
 #include "clang/CodeGen/SwiftCallingConv.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/Assumptions.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/DataLayout.h"
@@ -2015,6 +2016,18 @@ void CodeGenModule::ConstructAttributeList(
                                llvm::toStringRef(CodeGenOpts.UniformWGSize));
       }
     }
+
+    std::string AssumptionValueStr;
+    for (AssumptionAttr *AssumptionA :
+         TargetDecl->specific_attrs<AssumptionAttr>()) {
+      std::string AS = AssumptionA->getAssumption().str();
+      if (!AS.empty() && !AssumptionValueStr.empty())
+        AssumptionValueStr += ",";
+      AssumptionValueStr += AS;
+    }
+
+    if (!AssumptionValueStr.empty())
+      FuncAttrs.addAttribute(llvm::AssumptionAttrKey, AssumptionValueStr);
   }
 
   // Attach "no-builtins" attributes to:

clang/lib/Sema/CMakeLists.txt

@@ -1,4 +1,5 @@
 set(LLVM_LINK_COMPONENTS
+  Core
   FrontendOpenMP
   Support
   )

clang/lib/Sema/SemaDeclAttr.cpp

@@ -23,6 +23,7 @@
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/AST/Type.h"
 #include "clang/Basic/CharInfo.h"
+#include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/TargetBuiltins.h"
 #include "clang/Basic/TargetInfo.h"
@@ -38,6 +39,7 @@
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/IR/Assumptions.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -1712,6 +1714,42 @@ void Sema::AddAllocAlignAttr(Decl *D, const AttributeCommonInfo &CI,
   D->addAttr(::new (Context) AllocAlignAttr(Context, CI, Idx));
 }
 
+/// Check if \p AssumptionStr is a known assumption and warn if not.
+static void checkAssumptionAttr(Sema &S, SourceLocation Loc,
+                                StringRef AssumptionStr) {
+  if (llvm::KnownAssumptionStrings.count(AssumptionStr))
+    return;
+
+  unsigned BestEditDistance = 3;
+  StringRef Suggestion;
+  for (const auto &KnownAssumptionIt : llvm::KnownAssumptionStrings) {
+    unsigned EditDistance =
+        AssumptionStr.edit_distance(KnownAssumptionIt.getKey());
+    if (EditDistance < BestEditDistance) {
+      Suggestion = KnownAssumptionIt.getKey();
+      BestEditDistance = EditDistance;
+    }
+  }
+
+  if (!Suggestion.empty())
+    S.Diag(Loc, diag::warn_assume_attribute_string_unknown_suggested)
+        << AssumptionStr << Suggestion;
+  else
+    S.Diag(Loc, diag::warn_assume_attribute_string_unknown) << AssumptionStr;
+}
+
+static void handleAssumumptionAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
+  // Handle the case where the attribute has a text message.
+  StringRef Str;
+  SourceLocation AttrStrLoc;
+  if (!S.checkStringLiteralArgumentAttr(AL, 0, Str, &AttrStrLoc))
+    return;
+
+  checkAssumptionAttr(S, AttrStrLoc, Str);
+
+  D->addAttr(::new (S.Context) AssumptionAttr(S.Context, AL, Str));
+}
+
 /// Normalize the attribute, __foo__ becomes foo.
 /// Returns true if normalization was applied.
 static bool normalizeName(StringRef &AttrName) {
@@ -7845,6 +7883,9 @@ static void ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D,
   case ParsedAttr::AT_Unavailable:
     handleAttrWithMessage<UnavailableAttr>(S, D, AL);
     break;
+  case ParsedAttr::AT_Assumption:
+    handleAssumumptionAttr(S, D, AL);
+    break;
   case ParsedAttr::AT_ObjCDirect:
     handleObjCDirectAttr(S, D, AL);
     break;

clang/test/CodeGen/assume_attr.c

@@ -0,0 +1,58 @@
+// RUN: %clang_cc1 -emit-llvm -triple i386-linux-gnu %s -o - | FileCheck %s
+// RUN: %clang_cc1 -x c -emit-pch -o %t %s
+// RUN: %clang_cc1 -include-pch %t %s -emit-llvm -o - | FileCheck %s
+
+// TODO: for "foo" and "bar", "after" is not added as it appears "after" the first use or definition respectively. There might be a way to allow that.
+
+// CHECK:   define{{.*}} void @bar() #0
+// CHECK:   define{{.*}} void @baz() #1
+// CHECK:   declare{{.*}} void @foo() #2
+// CHECK:      attributes #0
+// CHECK-SAME:   "llvm.assume"="bar:before1,bar:before2,bar:before3,bar:def1,bar:def2"
+// CHECK:      attributes #1
+// CHECK-SAME:   "llvm.assume"="baz:before1,baz:before2,baz:before3,baz:def1,baz:def2,baz:after"
+// CHECK:      attributes #2
+// CHECK-SAME:   "llvm.assume"="foo:before1,foo:before2,foo:before3"
+
+#ifndef HEADER
+#define HEADER
+
+/// foo: declarations only
+
+__attribute__((assume("foo:before1"))) void foo(void);
+
+__attribute__((assume("foo:before2")))
+__attribute__((assume("foo:before3"))) void
+foo(void);
+
+/// baz: static function declarations and a definition
+
+__attribute__((assume("baz:before1"))) static void baz(void);
+
+__attribute__((assume("baz:before2")))
+__attribute__((assume("baz:before3"))) static void
+baz(void);
+
+// Definition
+__attribute__((assume("baz:def1,baz:def2"))) static void baz(void) { foo(); }
+
+__attribute__((assume("baz:after"))) static void baz(void);
+
+/// bar: external function declarations and a definition
+
+__attribute__((assume("bar:before1"))) void bar(void);
+
+__attribute__((assume("bar:before2")))
+__attribute__((assume("bar:before3"))) void
+bar(void);
+
+// Definition
+__attribute__((assume("bar:def1,bar:def2"))) void bar(void) { baz(); }
+
+__attribute__((assume("bar:after"))) void bar(void);
+
+/// back to foo
+
+__attribute__((assume("foo:after"))) void foo(void);
+
+#endif

clang/test/CodeGenCXX/assume_attr.cpp

@@ -0,0 +1,120 @@
+// RUN: %clang_cc1 -emit-llvm -triple i386-linux-gnu %s -o - | FileCheck %s
+// RUN: %clang_cc1 -x c++ -emit-pch -triple i386-linux-gnu -o %t %s
+// RUN: %clang_cc1 -include-pch %t %s -triple i386-linux-gnu -emit-llvm -o - | FileCheck %s
+// expected-no-diagnostics
+
+#ifndef HEADER
+#define HEADER
+
+/// foo: declarations only
+
+__attribute__((assume("foo:before1"))) void foo();
+
+__attribute__((assume("foo:before2")))
+__attribute__((assume("foo:before3"))) void
+foo();
+
+/// baz: static function declarations and a definition
+
+__attribute__((assume("baz:before1"))) static void baz();
+
+__attribute__((assume("baz:before2")))
+__attribute__((assume("baz:before3"))) static void
+baz();
+
+// Definition
+__attribute__((assume("baz:def1,baz:def2"))) static void baz() { foo(); }
+
+__attribute__((assume("baz:after"))) static void baz();
+
+/// bar: external function declarations and a definition
+
+__attribute__((assume("bar:before1"))) void bar();
+
+__attribute__((assume("bar:before2")))
+__attribute__((assume("bar:before3"))) void
+bar();
+
+// Definition
+__attribute__((assume("bar:def1,bar:def2"))) void bar() { baz(); }
+
+__attribute__((assume("bar:after"))) void bar();
+
+/// back to foo
+
+__attribute__((assume("foo:after"))) void foo();
+
+/// class tests
+class C {
+  __attribute__((assume("C:private_method"))) void private_method();
+  __attribute__((assume("C:private_static"))) static void private_static();
+
+public:
+  __attribute__((assume("C:public_method1"))) void public_method();
+  __attribute__((assume("C:public_static1"))) static void public_static();
+};
+
+__attribute__((assume("C:public_method2"))) void C::public_method() {
+  private_method();
+}
+
+__attribute__((assume("C:public_static2"))) void C::public_static() {
+  private_static();
+}
+
+/// template tests
+template <typename T>
+__attribute__((assume("template_func<T>"))) void template_func() {}
+
+template <>
+__attribute__((assume("template_func<float>"))) void template_func<float>() {}
+
+template <>
+void template_func<int>() {}
+
+template <typename T>
+struct S {
+  __attribute__((assume("S<T>::method"))) void method();
+};
+
+template <>
+__attribute__((assume("S<float>::method"))) void S<float>::method() {}
+
+template <>
+void S<int>::method() {}
+
+// CHECK:         define{{.*}} void @_Z3barv() #0
+// CHECK:         define{{.*}} void @_ZL3bazv() #1
+// CHECK:         define{{.*}} void @_ZN1C13public_methodEv({{.*}}) #2
+// CHECK:         declare{{.*}} void @_ZN1C14private_methodEv({{.*}}) #3
+// CHECK:         define{{.*}} void @_ZN1C13public_staticEv() #4
+// CHECK:         declare{{.*}} void @_ZN1C14private_staticEv() #5
+// CHECK:         define{{.*}} void @_Z13template_funcIfEvv() #6
+// CHECK:         define{{.*}} void @_Z13template_funcIiEvv() #7
+// CHECK:         define{{.*}} void @_ZN1SIfE6methodEv({{.*}}) #8
+// CHECK:         define{{.*}} void @_ZN1SIiE6methodEv({{.*}}) #9
+// CHECK:         declare{{.*}} void @_Z3foov() #10
+// CHECK:         attributes #0
+// CHECK-SAME:      "llvm.assume"="bar:before1,bar:before2,bar:before3,bar:def1,bar:def2"
+// CHECK:         attributes #1
+// CHECK-SAME:      "llvm.assume"="baz:before1,baz:before2,baz:before3,baz:def1,baz:def2,baz:after"
+// CHECK:         attributes #2
+// CHECK-SAME:      "llvm.assume"="C:public_method1,C:public_method2"
+// CHECK:         attributes #3
+// CHECK-SAME:      "llvm.assume"="C:private_method"
+// CHECK:         attributes #4
+// CHECK-SAME:      "llvm.assume"="C:public_static1,C:public_static2"
+// CHECK:         attributes #5
+// CHECK-SAME:      "llvm.assume"="C:private_static"
+// CHECK:         attributes #6
+// CHECK-SAME:      "llvm.assume"="template_func<T>,template_func<float>"
+// CHECK:         attributes #7
+// CHECK-SAME:      "llvm.assume"="template_func<T>"
+// CHECK:         attributes #8
+// CHECK-SAME:      "llvm.assume"="S<T>::method,S<float>::method"
+// CHECK:         attributes #9
+// CHECK-SAME:      "llvm.assume"="S<T>::method"
+// CHECK:         attributes #10
+// CHECK-SAME:      "llvm.assume"="foo:before1,foo:before2,foo:before3"
+
+#endif

clang/test/Misc/pragma-attribute-supported-attributes-list.test

@@ -20,6 +20,7 @@
 // CHECK-NEXT: ArcWeakrefUnavailable (SubjectMatchRule_objc_interface)
 // CHECK-NEXT: ArmBuiltinAlias (SubjectMatchRule_function)
 // CHECK-NEXT: AssumeAligned (SubjectMatchRule_objc_method, SubjectMatchRule_function)
+// CHECK-NEXT: Assumption (SubjectMatchRule_function SubjectMatchRule_objc_method)
 // CHECK-NEXT: Availability ((SubjectMatchRule_record, SubjectMatchRule_enum, SubjectMatchRule_enum_constant, SubjectMatchRule_field, SubjectMatchRule_function, SubjectMatchRule_namespace, SubjectMatchRule_objc_category, SubjectMatchRule_objc_implementation, SubjectMatchRule_objc_interface, SubjectMatchRule_objc_method, SubjectMatchRule_objc_property, SubjectMatchRule_objc_protocol, SubjectMatchRule_record, SubjectMatchRule_type_alias, SubjectMatchRule_variable))
 // CHECK-NEXT: BPFPreserveAccessIndex (SubjectMatchRule_record)
 // CHECK-NEXT: CFAuditedTransfer (SubjectMatchRule_function)

clang/test/Sema/attr-assume.c

@@ -0,0 +1,14 @@
+// RUN: %clang_cc1 -triple i386-apple-darwin9 -fsyntax-only -verify %s
+
+void f1() __attribute__((assume(3))); // expected-error {{'assume' attribute requires a string}}
+void f2() __attribute__((assume(int))); // expected-error {{expected expression}}
+void f3() __attribute__((assume(for))); // expected-error {{expected expression}}
+void f4() __attribute__((assume("QQQQ"))); // expected-warning {{unknown assumption string 'QQQQ'; attribute is potentially ignored}}
+void f5() __attribute__((assume("omp_no_openmp")));
+void f6() __attribute__((assume("omp_noopenmp"))); // expected-warning {{unknown assumption string 'omp_noopenmp' may be misspelled; attribute is potentially ignored, did you mean 'omp_no_openmp'?}}
+void f7() __attribute__((assume("omp_no_openmp_routine"))); // expected-warning {{unknown assumption string 'omp_no_openmp_routine' may be misspelled; attribute is potentially ignored, did you mean 'omp_no_openmp_routines'?}}
+void f8() __attribute__((assume("omp_no_openmp1"))); // expected-warning {{unknown assumption string 'omp_no_openmp1' may be misspelled; attribute is potentially ignored, did you mean 'omp_no_openmp'?}}
+void f9() __attribute__((assume("omp_no_openmp", "omp_no_openmp"))); // expected-error {{'assume' attribute takes one argument}}
+
+int g1 __attribute__((assume(0))); // expected-warning {{'assume' attribute only applies to functions and Objective-C methods}}
+int g2 __attribute__((assume("omp_no_openmp"))); // expected-warning {{'assume' attribute only applies to functions and Objective-C methods}}

llvm/include/llvm/IR/Assumptions.h

@@ -0,0 +1,50 @@
+//===--- Assumptions.h - Assumption handling and organization ---*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// String assumptions that are known to optimization passes should be placed in
+// the KnownAssumptionStrings set. This can be done in various ways, i.a.,
+// via a static KnownAssumptionString object.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_IR_ASSUMPTIONS_H
+#define LLVM_IR_ASSUMPTIONS_H
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
+
+namespace llvm {
+
+class Function;
+
+/// The key we use for assumption attributes.
+constexpr StringRef AssumptionAttrKey = "llvm.assume";
+
+/// A set of known assumption strings that are accepted without warning and
+/// which can be recommended as typo correction.
+extern StringSet<> KnownAssumptionStrings;
+
+/// Helper that allows to insert a new assumption string in the known assumption
+/// set by creating a (static) object.
+struct KnownAssumptionString {
+  KnownAssumptionString(StringRef AssumptionStr)
+      : AssumptionStr(AssumptionStr) {
+    KnownAssumptionStrings.insert(AssumptionStr);
+  }
+  operator StringRef() const { return AssumptionStr; }
+
+private:
+  StringRef AssumptionStr;
+};
+
+/// Return true if \p F has the assumption \p AssumptionStr attached.
+bool hasAssumption(Function &F, const KnownAssumptionString &AssumptionStr);
+
+} // namespace llvm
+
+#endif

llvm/lib/IR/Assumptions.cpp

@@ -0,0 +1,36 @@
+//===- Assumptions.cpp ------ Collection of helpers for assumptions -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/IR/Assumptions.h"
+#include "llvm/IR/Attributes.h"
+#include "llvm/IR/Function.h"
+
+using namespace llvm;
+
+bool llvm::hasAssumption(Function &F,
+                         const KnownAssumptionString &AssumptionStr) {
+  const Attribute &A = F.getFnAttribute(AssumptionAttrKey);
+  if (!A.isValid())
+    return false;
+  assert(A.isStringAttribute() && "Expected a string attribute!");
+
+  SmallVector<StringRef, 8> Strings;
+  A.getValueAsString().split(Strings, ",");
+
+  return llvm::any_of(Strings, [=](StringRef Assumption) {
+    return Assumption == AssumptionStr;
+  });
+}
+
+StringSet<> llvm::KnownAssumptionStrings({
+    "omp_no_openmp",          // OpenMP 5.1
+    "omp_no_openmp_routines", // OpenMP 5.1
+    "omp_no_parallelism",     // OpenMP 5.1
+});

llvm/lib/IR/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_llvm_component_library(LLVMCore
   AbstractCallSite.cpp
   AsmWriter.cpp
+  Assumptions.cpp
   Attributes.cpp
   AutoUpgrade.cpp
   BasicBlock.cpp