address review comment

Signed-off-by: Whitney Tsang <whitney.tsang@intel.com>
intel · May 20, 2024 · 64f05e8 · 64f05e8
1 parent a4eed22
commit 64f05e8
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/third_party/intel/lib/Target/LLVMIR/Dialect/TritonGEN/TritonGENToLLVMIRTranslation.cpp b/third_party/intel/lib/Target/LLVMIR/Dialect/TritonGEN/TritonGENToLLVMIRTranslation.cpp
@@ -35,6 +35,10 @@ class TritonGENDialectLLVMIRTranslationInterface
   amendOperation(Operation *op, ArrayRef<llvm::Instruction *> instructions,
                  NamedAttribute attribute,
                  LLVM::ModuleTranslation &moduleTranslation) const final {
+    // Skip the attribute if it is not a TritonGEN attribute.
+    if (!attribute.getName().getValue().starts_with("triton_gen"))
+      return success();
+
     llvm::LLVMContext &llvmContext = moduleTranslation.getLLVMContext();
     llvm::Function *llvmFunc =
         moduleTranslation.lookupFunction(cast<LLVM::LLVMFuncOp>(op).getName());
@@ -53,6 +57,14 @@ class TritonGENDialectLLVMIRTranslationInterface
   // The attribute is converted into metadata and added to the function.
   void amendKernel(llvm::LLVMContext &llvmContext, llvm::Function *llvmFunc,
                    NamedAttribute attribute) const {
+    StringRef name = attribute.getName().getValue();
+    assert((name == triton::TritonGEN::TritonGENDialect::
+                        getMaxWorkGroupSizeAttrName() ||
+            name == triton::TritonGEN::TritonGENDialect::
+                        getReqdWorkGroupSizeAttrName() ||
+            name == triton::TritonGEN::TritonGENDialect::
+                        getReqdSubGroupSizeAttrName()) &&
+           "Unexpected attribute");
     SmallVector<llvm::Metadata *, 3> metadata;
     llvm::Type *i64 = llvm::IntegerType::get(llvmContext, 64);
     for (int64_t i :
@@ -61,7 +73,7 @@ class TritonGENDialectLLVMIRTranslationInterface
       metadata.push_back(llvm::ConstantAsMetadata::get(constant));
     }
     llvm::MDNode *node = llvm::MDNode::get(llvmContext, metadata);
-    llvmFunc->setMetadata(attribute.getName().getValue().drop_front(11), node);
+    llvmFunc->setMetadata(name.drop_front(11), node);
   }
 };
 } // namespace