From dd03f957ecab48024d01fc84f16e4be7e1447cf6 Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <kparzysz@quicinc.com>
Date: Mon, 26 Oct 2020 10:01:02 -0500
Subject: [PATCH] [LLVM] Add target feature string to function attributes

---
 src/target/llvm/codegen_llvm.cc                  |  5 +++++
 .../unittest/test_target_codegen_hexagon.py      | 16 ++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 2a7e4644571b..fd55f2418628 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -172,6 +172,11 @@ void CodeGenLLVM::AddFunctionInternal(const PrimFunc& f, bool ret_void) {
   }
 #endif
 
+  llvm::StringRef fs = target_machine_->getTargetFeatureString();
+  if (!fs.empty()) {
+    function_->addFnAttr("target-features", fs);
+  }
+
   if (ret_void) {
     builder_->CreateRetVoid();
   } else {
diff --git a/tests/python/unittest/test_target_codegen_hexagon.py b/tests/python/unittest/test_target_codegen_hexagon.py
index d42693ef16bd..b74d487f3fa7 100644
--- a/tests/python/unittest/test_target_codegen_hexagon.py
+++ b/tests/python/unittest/test_target_codegen_hexagon.py
@@ -63,6 +63,21 @@ def check_add(offload):
     check_add(False)
 
 
+def test_llvm_target_features():
+    if not check_prereq_and_setup():
+        return
+    target = tvm.target.hexagon("v66", hvx=128)
+    # Define some trivial compute
+    A = tvm.te.placeholder((128,), dtype="uint8", name="A")
+    C = tvm.te.compute((128,), lambda i: A[i] + 1, name="C")
+    s = tvm.te.create_schedule(C.op)
+    m = tvm.build(s, [C, A], target=target, target_host=target, name="add_one")
+    llvm_ir = m.get_source("ll")
+    # Make sure we find +hvx-length128b in "attributes".
+    fs = re.findall(r"attributes.*\+hvx-length128b", llvm_ir)
+    assert fs  # Check that it's non-empty
+
+
 def test_alloc_vtcm():
     if not check_prereq_and_setup():
         return
@@ -92,4 +107,5 @@ def test_alloc_vtcm():
 
 if __name__ == "__main__":
     test_basic()
+    test_llvm_target_features()
     test_alloc_vtcm()