[AMD] Correctly mask program_id(1) for architected SGPRs (#7455)

Architectures with the architected SGPRs feature (Just gfx12 for now)
don't use sX registers for the program id,
but read from registers ttmp9 (X) and ttmp7(Y[15:0], Z[31:16]). As a
result, the upper 16 bits of ttmp7 need to be masked when reading Y.
This can be omitted when Z is not used. optimize_module infers this from
existence of @llvm.amdgcn.workgroup.id.z() and creates the attribute
`amdgpu-no-workgroup-id-z` for the backend.
`test_core.py:test_zero_strided_tensors` doesn't read Z, but still has
thread blocks on that axis. I don't think we know for sure during
compilation which dispatch dimensions are going to be used. We can query
this feature directly from LLVM and just set all dimensions as used.
This does not increase register pressure, because these registers are
reserved anyway. It only adds an `s_and_b32 s1, ttmp7, 0xffff`
instruction when loading Y.
Not too happy with recreating the MCSubtargetInfo just for this, but I
didn't see a better way to preserve this object. With this change,
`test_core.py:test_zero_strided_tensors` passes on gfx12.

See "3.5.3.4. Compute Shader (CS)" in
https://www.amd.com/content/dam/amd/en/documents/radeon-tech-docs/instruction-set-architectures/rdna4-instruction-set-architecture.pdf
for details on this.

---------

Co-authored-by: Paul Trojahn <[email protected]>

Author: ptrojahn

python/src/llvm.cc

@@ -220,6 +220,8 @@ void init_triton_llvm(py::module &&m) {
       .def("set_calling_conv", &llvm::Function::setCallingConv)
       .def("add_fn_attr", [](llvm::Function *fn, std::string &name,
                              std::string &val) { fn->addFnAttr(name, val); })
+      .def("remove_fn_attr", [](llvm::Function *fn,
+                                std::string &name) { fn->removeFnAttr(name); })
       .def("add_fn_asan_attr",
            [](llvm::Function *fn) {
              fn->addFnAttr(llvm::Attribute::SanitizeAddress);

third_party/amd/backend/compiler.py

@@ -377,6 +377,15 @@ def make_llir(src, metadata, options):
 
         llvm.optimize_module(llvm_mod, llvm.OPTIMIZE_O3, options.arch, '', [], options.enable_fp_fusion)
 
+        # Architectures with architected SGPRs store the workgroup id in ttmp9 (X) and ttmp7 (Y[15:0], Z[31:16]).
+        # These attributes are used to determine if Z should be masked out when loading Y. They are inferred during
+        # optimize_module from calls to @llvm.amdgcn.workgroup.id.x/y/z(). We cannot rely on this because a
+        # dispatch dimensions might be used even if there is no program_id() call for it.
+        if amd.has_architected_sgprs(options.arch):
+            fns[0].remove_fn_attr("amdgpu-no-workgroup-id-x")
+            fns[0].remove_fn_attr("amdgpu-no-workgroup-id-y")
+            fns[0].remove_fn_attr("amdgpu-no-workgroup-id-z")
+
         if knobs.amd.scalarize_packed_fops:
             amd.add_scalarize_packed_fops_llvm_pass(fns[0])
 

third_party/amd/python/triton_amd.cc

@@ -261,6 +261,18 @@ void init_triton_amd(py::module &&m) {
       },
       py::return_value_policy::take_ownership);
 
+  m.def("has_architected_sgprs", [](const std::string &arch) {
+    std::string error;
+    llvm::Triple triple(amdTargetTriple);
+    const llvm::Target *target =
+        llvm::TargetRegistry::lookupTarget(triple.normalize(), error);
+    if (!target)
+      throw std::runtime_error("target lookup error: " + error);
+    std::unique_ptr<llvm::MCSubtargetInfo> sti(
+        target->createMCSubtargetInfo(amdTargetTriple, arch, ""));
+    return sti->checkFeatures("+architected-sgprs");
+  });
+
   m.def("need_extern_lib", [](llvm::Module *module, const std::string &lib) {
     for (llvm::Function &f : module->functions()) {
       if (f.hasExternalLinkage() && f.hasName() && !f.hasExactDefinition()) {