Merge commit '04f87d021a3550aa536862aecee21bf0a30a2452'

Author: whitneywhtsang

python/test/unit/language/test_libdevice.py

@@ -4,6 +4,7 @@
 import triton
 import triton.language as tl
 from triton.language.extra.intel import libdevice
+from triton.language.extra.libdevice import fast_dividef as my_fast_dividef
 
 
 @pytest.mark.parametrize("dtype_str", ["float32", "float64"])
@@ -39,3 +40,20 @@ def kernel(in_p, out_p, fn: tl.constexpr, SIZE: tl.constexpr):
     kernel[(1, )](x, y_exp, fn=libdevice_fn, SIZE=SIZE, num_warps=4, num_ctas=1)
 
     torch.testing.assert_close(y_ref, y_exp, equal_nan=True)
+
+
+def test_libdevice_rename(device):
+    # mark the import as used by this test
+    _ = my_fast_dividef
+
+    @triton.jit
+    def triton_copy(in_ptr, out_ptr, BLOCK_SIZE: tl.constexpr):
+        offsets = tl.arange(0, BLOCK_SIZE)
+        data = tl.load(in_ptr + offsets)
+        tl.store(out_ptr + offsets, data)
+
+    BLOCK_SIZE = 256
+    inp = torch.randn(BLOCK_SIZE, device=device)
+    out = torch.empty_like(inp)
+
+    triton_copy[(1, )](inp, out, BLOCK_SIZE)

python/triton/compiler/code_generator.py

@@ -218,7 +218,7 @@ def __init__(self, context, prototype, gscope, attributes, constants, function_n
 
             module_name = getattr(v, "__module__", "")
             if module_name in module_map:
-                self.gscope[k] = getattr(module_map[module_name], k)
+                self.gscope[k] = getattr(module_map[module_name], v.__name__)
             else:
                 self.gscope[k] = v