1 files changed, 17 insertions, 3 deletions
diff --git a/lib/Target/R600/SIMachineFunctionInfo.cpp b/lib/Target/R600/SIMachineFunctionInfo.cpp
index 7e59b42..1a4e4cb 100644
--- a/lib/Target/R600/SIMachineFunctionInfo.cpp
+++ b/lib/Target/R600/SIMachineFunctionInfo.cpp
@@ -10,11 +10,25 @@
 
 
 #include "SIMachineFunctionInfo.h"
+#include "llvm/IR/Attributes.h"
+#include "llvm/IR/Function.h"
 
 using namespace llvm;
 
+const char *SIMachineFunctionInfo::ShaderTypeAttribute = "ShaderType";
+
 SIMachineFunctionInfo::SIMachineFunctionInfo(const MachineFunction &MF)
   : MachineFunctionInfo(),
-    SPIPSInputAddr(0),
-    ShaderType(0)
-  { }
+    ShaderType(0),
+    PSInputAddr(0) {
+
+  AttributeSet Set = MF.getFunction()->getAttributes();
+  Attribute A = Set.getAttribute(AttributeSet::FunctionIndex,
+                                 ShaderTypeAttribute);
+
+  if (A.isStringAttribute()) {
+    StringRef Str = A.getValueAsString();
+    if (Str.getAsInteger(0, ShaderType))
+      llvm_unreachable("Can't parse shader type!");
+  }
+}