Fix indeterminism bug in LLVM datalayout (#6303)

The data layout depends on the flag `nvptx-short-ptr` but this flag was
being set after we set the data layout. Since this is a global this was
only affecting the first kernel compiled.

Author: ThomasRaoux

third_party/nvidia/backend/compiler.py

@@ -331,6 +331,7 @@ def make_llir(self, src, metadata, options, capability):
         proc = sm_arch_from_capability(capability)
         features = get_features(options, self.target.arch)
         triple = 'nvptx64-nvidia-cuda'
+        nvidia.set_short_ptr()
         llvm.attach_datalayout(llvm_mod, triple, proc, features)
         nvidia.set_nvvm_reflect_ftz(llvm_mod)
 
@@ -366,7 +367,7 @@ def make_ptx(self, src, metadata, opt, capability):
         triple = 'nvptx64-nvidia-cuda'
         proc = sm_arch_from_capability(capability)
         features = get_features(opt, self.target.arch)
-        ret = llvm.translate_to_asm(src, triple, proc, features, ['nvptx-short-ptr'], opt.enable_fp_fusion, False)
+        ret = llvm.translate_to_asm(src, triple, proc, features, [], opt.enable_fp_fusion, False)
         # Find kernel names (there should only be one)
         names = re.findall(r".visible .entry ([a-zA-Z_][a-zA-Z0-9_]*)", ret)
         assert len(names) == 1

third_party/nvidia/triton_nvidia.cc

@@ -79,6 +79,16 @@ void init_triton_nvidia(py::module &&m) {
     context.loadAllAvailableDialects();
   });
 
+  // Set short point option, this needs to be set before setting the data
+  // layout.
+  m.def("set_short_ptr", []() {
+    auto options = llvm::cl::getRegisteredOptions();
+    const char *flag = "nvptx-short-ptr";
+    auto *shortPtr = static_cast<llvm::cl::opt<bool> *>(options[flag]);
+    assert(shortPtr);
+    shortPtr->setValue(true);
+  });
+
   // TODO: could be done in python if we had a generic interface to set metadata
   m.def("set_nvvm_reflect_ftz", [](llvm::Module *mod) {
     // please check https://llvm.org/docs/NVPTXUsage.html#reflection-parameters