Add more libdevice lowerings (#97)

triton-lang · Aug 9, 2024 · 5662d4c · 5662d4c
1 parent 9035f4b
commit 5662d4c
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 3 deletions.
diff --git a/python/src/ir.cc b/python/src/ir.cc
@@ -1,4 +1,4 @@
-#include <pybind11/functional.h>
+#include <pybind11/functional.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
@@ -1425,6 +1425,10 @@ void init_triton_ir(py::module &&m) {
            [](TritonOpBuilder &self, Value &val) -> Value {
              return self.create<math::Exp2Op>(val);
            })
+      .def("create_expm1",
+           [](TritonOpBuilder &self, Value &val) -> Value {
+             return self.create<math::ExpM1Op>(val);
+           })
       .def("create_cos",
            [](TritonOpBuilder &self, Value &val) -> Value {
              return self.create<math::CosOp>(val);
@@ -1477,6 +1481,10 @@ void init_triton_ir(py::module &&m) {
            [](TritonOpBuilder &self, Value &val) -> Value {
              return self.create<math::LogOp>(val);
            })
+      .def("create_log1p",
+           [](TritonOpBuilder &self, Value &val) -> Value {
+             return self.create<math::Log1pOp>(val);
+           })
       .def("create_log2",
            [](TritonOpBuilder &self, Value &val) -> Value {
              return self.create<math::Log2Op>(val);

diff --git a/python/test/unit/cpu/test_libdevice.py b/python/test/unit/cpu/test_libdevice.py
@@ -21,14 +21,18 @@ def is_cpu():
 
 @pytest.mark.parametrize("dtype_str", float_dtypes)
 @pytest.mark.parametrize("math_fn", [
-    "acos", "acosh", "asin", "asinh", "atan", "atanh", "cbrt", "cos", "cosh", "erf", "exp", "exp2", "log", "log2",
-    "log10", "sin", "sinh", "tan", "tanh"
+    "acos", "acosh", "asin", "asinh", "atan", "atanh", "cbrt", "cos", "cosh", "erf", "exp", "exp2", "expm1", "floor",
+    "log", "log1p", "log2", "log10", "rsqrt", "sin", "sinh", "sqrt", "tan", "tanh"
 ])
 @pytest.mark.parametrize("size", [1, 4, 16, 64])
 def test_libdevice(dtype_str, math_fn, size, device):
     if not is_cpu():
         pytest.skip("This test is CPU-specific")
 
+    if dtype_str == "bfloat16":
+        if math_fn == "floor" or math_fn == "rsqrt":
+            pytest.skip("libgcc < 13 does not define __truncsfbf2, which this op needs")
+
     @triton.jit
     def kernel(src, dst, MATH_FN: tl.constexpr, BLOCK_SIZE: tl.constexpr):
         idxs = tl.arange(0, BLOCK_SIZE)

diff --git a/python/triton/language/extra/cpu/libdevice.py b/python/triton/language/extra/cpu/libdevice.py
@@ -61,6 +61,16 @@ def exp2(arg0, _builder=None):
     return core.tensor(_builder.create_exp2(arg0.handle), arg0.type)
 
 
+@core.extern
+def expm1(arg0, _builder=None):
+    return core.tensor(_builder.create_expm1(arg0.handle), arg0.type)
+
+
+@core.extern
+def floor(arg0, _builder=None):
+    return core.tensor(_builder.create_floor(arg0.handle), arg0.type)
+
+
 @core.extern
 def log(arg0, _builder=None):
     return core.tensor(_builder.create_log(arg0.handle), arg0.type)
@@ -76,11 +86,26 @@ def log10(arg0, _builder=None):
     return core.tensor(_builder.create_log10(arg0.handle), arg0.type)
 
 
+@core.extern
+def log1p(arg0, _builder=None):
+    return core.tensor(_builder.create_log1p(arg0.handle), arg0.type)
+
+
 @core.extern
 def sin(arg0, _builder=None):
     return core.tensor(_builder.create_sin(arg0.handle), arg0.type)
 
 
+@core.extern
+def rsqrt(arg0, _builder=None):
+    return core.tensor(_builder.create_rsqrt(arg0.handle), arg0.type)
+
+
+@core.extern
+def sqrt(arg0, _builder=None):
+    return core.tensor(_builder.create_sqrt(arg0.handle), arg0.type)
+
+
 @core.extern
 def sinh(arg0, _builder=None):
     return core.tensor(_builder.create_sinh(arg0.handle), arg0.type)

diff --git a/third_party/cpu/lib/TritonToTritonCPU/ConvertElementwiseOps.cpp b/third_party/cpu/lib/TritonToTritonCPU/ConvertElementwiseOps.cpp
@@ -220,9 +220,11 @@ struct ConvertElementwiseOps
     patterns.add<OpTypeConversion<math::AbsIOp>>(typeConverter, context);
     patterns.add<OpTypeConversion<math::ExpOp>>(typeConverter, context);
     patterns.add<OpTypeConversion<math::Exp2Op>>(typeConverter, context);
+    patterns.add<OpTypeConversion<math::ExpM1Op>>(typeConverter, context);
     patterns.add<OpTypeConversion<math::LogOp>>(typeConverter, context);
     patterns.add<OpTypeConversion<math::Log2Op>>(typeConverter, context);
     patterns.add<OpTypeConversion<math::Log10Op>>(typeConverter, context);
+    patterns.add<OpTypeConversion<math::Log1pOp>>(typeConverter, context);
     patterns.add<OpTypeConversion<math::SinOp>>(typeConverter, context);
     patterns.add<OpTypeConversion<math::SinhOp>>(typeConverter, context);
     patterns.add<OpTypeConversion<math::CosOp>>(typeConverter, context);