From 44eb3891ae111fc8401d2a6717b0403cfc5a0c34 Mon Sep 17 00:00:00 2001
From: Philippe Tillet <ptillet@g.harvard.edu>
Date: Fri, 23 Aug 2019 20:29:12 -0700
Subject: [PATCH] [lang] added support for restrict; added macros for
 attributes

---
 examples/cpp/dot.cc       | 17 ++++++++++++-----
 include/triton/lang/ast.h | 11 ++++++++++-
 lib/lang/code_gen.cc      | 12 +++++++-----
 lib/lang/parser.cc        | 19 +++++++++++++++++--
 4 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/examples/cpp/dot.cc b/examples/cpp/dot.cc
index 3122ead6f..7ca9a1937 100644
--- a/examples/cpp/dot.cc
+++ b/examples/cpp/dot.cc
@@ -82,14 +82,21 @@ R"(
 #define true 1
 #define false 0
 #define __bool_true_false_are_defined 1
+
+#define __readonly      __attribute__((readonly))
+#define __writeonly     __attribute__((writeonly))
+#define __noalias       __attribute__((noalias))
+#define __aligned(A)    __attribute__((aligned(A)))
+#define __multipleof(A) __attribute__((multipleof(A)))
+
 extern int get_program_id(int);
 
-void matmul(restrict )" + a_ty + R"( * A __attribute__((readonly, aligned(16))),
-            restrict )" + b_ty + R"( * B __attribute__((readonly, aligned(16))),
-            restrict )" + c_ty + R"( * C __attribute__((aligned(16))),
+void matmul()" + a_ty + R"( * A __noalias __readonly __aligned(16),
+            )" + b_ty + R"( * B __noalias __readonly __aligned(16),
+            )" + c_ty + R"( * C __noalias __readonly __aligned(16),
             int M, int N, int K,
-            int lda __attribute__((multiple_of(8))),
-            int ldb __attribute__((multiple_of(8))),
+            int lda __multipleof(8),
+            int ldb __multipleof(8),
             int ldc) {
   int ridx = get_program_id(0);
   int ridy = get_program_id(1);
diff --git a/include/triton/lang/ast.h b/include/triton/lang/ast.h
index d293a0ae5..710c67e4a 100644
--- a/include/triton/lang/ast.h
+++ b/include/triton/lang/ast.h
@@ -58,7 +58,16 @@ class TranslationUnit;
 class ASTNode {
 public:
   struct Attr{
-    std::string name;
+
+    enum KindT{
+      MULTIPLEOF,
+      ALIGNED,
+      NOALIAS,
+      READONLY,
+      WRITEONLY
+    };
+
+    KindT kind;
     std::vector<Expr*> vals;
   };
   using AttrList = std::vector<Attr>;
diff --git a/lib/lang/code_gen.cc b/lib/lang/code_gen.cc
index f7dd021ed..451015d84 100644
--- a/lib/lang/code_gen.cc
+++ b/lib/lang/code_gen.cc
@@ -356,6 +356,8 @@ void Generator::VisitFuncDef(FuncDef* funcDef) {
     args[i]->set_name(name);
     for(ASTNode::Attr attr: obj->GetAttrList())
       fn->add_attr(i, GenIRAttr(attr));
+    if(obj->IsRestrictQualified())
+      fn->add_attr(i, ir::attribute(ir::noalias));
     mod_->set_value(name, nullptr, args[i]);
     mod_->get_scope().types[name] = args[i]->get_type();
     i++;
@@ -440,22 +442,22 @@ ir::value* Generator::GenCastOp(ir::value* src, ir::type* dst_ty) {
 
 // Triton-IR Attr
 ir::attribute Generator::GenIRAttr(ASTNode::Attr attr) {
-  if(attr.name == "multiple_of") {
+  if(attr.kind == ASTNode::Attr::MULTIPLEOF) {
     VisitExpr(attr.vals[0]);
     auto cst = dynamic_cast<ir::constant_int*>(ret_);
     if(!cst) should_not_happen();
     return ir::attribute(ir::multiple_of, cst->get_value());
   }
-  if(attr.name == "aligned") {
+  if(attr.kind == ASTNode::Attr::ALIGNED) {
     VisitExpr(attr.vals[0]);
     auto cst = dynamic_cast<ir::constant_int*>(ret_);
     return ir::attribute(ir::aligned, cst->get_value());
   }
-  if(attr.name == "noalias")
+  if(attr.kind == ASTNode::Attr::NOALIAS)
     return ir::attribute(ir::noalias);
-  if(attr.name == "readonly")
+  if(attr.kind == ASTNode::Attr::READONLY)
     return ir::attribute(ir::readonly);
-  if(attr.name == "writeonly")
+  if(attr.kind == ASTNode::Attr::WRITEONLY)
     return ir::attribute(ir::writeonly);
   should_not_happen();
 }
diff --git a/lib/lang/parser.cc b/lib/lang/parser.cc
index 35ed63e15..d8dbcf1c8 100644
--- a/lib/lang/parser.cc
+++ b/lib/lang/parser.cc
@@ -1806,7 +1806,8 @@ Object* Parser::ParseParamDecl() {
   auto type = ParseDeclSpec(&storageSpec, &funcSpec, nullptr);
   auto tokTypePair = ParseDeclarator(type);
   auto tok = tokTypePair.tok;
-  type = Type::MayCast(tokTypePair.type, true);
+  QualType fullType(tokTypePair.type.GetPtr(), type.Qual());
+  type = Type::MayCast(fullType, true);
   auto attrs = tokTypePair.attrs;
   if (!tok) { // Abstract declarator
     return Object::NewAnony(ts_.Peek(), type, 0, Linkage::L_NONE);
@@ -2692,7 +2693,21 @@ ASTNode::Attr Parser::ParseAttribute() {
   if (!ts_.Test(Token::IDENTIFIER))
     return ret;
   auto tok = ts_.Next();
-  ret.name = tok->str_;
+  std::string name = tok->str_;
+  // set kind
+  if(name == "aligned")
+    ret.kind = ASTNode::Attr::ALIGNED;
+  else if(name == "readonly")
+    ret.kind = ASTNode::Attr::READONLY;
+  else if(name == "writeonly")
+    ret.kind = ASTNode::Attr::WRITEONLY;
+  else if(name == "multipleof")
+    ret.kind = ASTNode::Attr::MULTIPLEOF;
+  else if(name == "noalias")
+    ret.kind = ASTNode::Attr::NOALIAS;
+  else
+    Error(tok, "unknown attribute kind");
+  // set exprs
   if (ts_.Try('(')) {
     if (ts_.Try(')'))
       return ret;