Update and cleanup FeStiff/Pythran

FeStiff/Pythran/main.py

@@ -1,54 +1,66 @@
-from rando import *
-import numpy as np
-from Stiffness import *
-from RandomTriangle import *
-import time
+from time import perf_counter as time
 import socket
 
+import numpy as np
+
+from rando import Rando
+from stiffness import Stiffness
+from random_triangle import random_triangle
+
+try:
+    from time import perf_counter_ns as time
+except ImportError:
+    pass
 
-ntri=10**7
-x=np.empty(3)
-y=np.empty(3)
-mat=np.empty(21)
-S=Stiffness()
+
+x = np.empty(3)
+y = np.empty(3)
+mat = np.empty(21)
+stiffness = Stiffness()
 
 print("\nVerify that, on the reference element, we are coherent with sage")
 print("(see ../sage/):\n")
-x[0]=0.0; x[1]=1.0; x[2]=0.0;
-y[0]=0.0; y[1]=0.0; y[2]=1.0;
-S.op(x,y,mat)
-
+x[0] = 0.0
+x[1] = 1.0
+x[2] = 0.0
+y[0] = 0.0
+y[1] = 0.0
+y[2] = 1.0
+stiffness.op(x, y, mat)
 
-for i in range(0,6):
-    print([mat[i*(i+1)//2+j] for j in range(0,i+1)])
+for i in range(6):
+    print([mat[i * (i + 1) // 2 + j] for j in range(0, i + 1)])
 print("\nWe must get the same result if we dilate the triangle:\n")
-for i in range(0,3):
-    x[i]*=2.
-    y[i]*=2
-S.op(x,y,mat)
-for i in range(0,6):
-    print([mat[i*(i+1)//2+j] for j in range(0,i+1)])
+for i in range(3):
+    x[i] *= 2.0
+    y[i] *= 2
+stiffness.op(x, y, mat)
+for i in range(6):
+    print([mat[i * (i + 1) // 2 + j] for j in range(0, i + 1)])
 
 print("\nNow, start the benchmark:")
-ntri=1000000
-print(ntri," triangles.")
-R = rando()
-t1 = time.time()
-for tr in range(0,ntri):
-    RandomTriangle(R,x,y)
-    S.op(x,y,mat)
-t=(time.time()-t1)
-print("first phase: ",t," seconds.")
-t1 = time.time()
-for tr in range(0,ntri):
-    RandomTriangle(R,x,y)
-tr=(time.time()-t1)
-print("second phase: ",tr," seconds.")
-
-t-=tr
-print("Total time: ",t," seconds.")
-print("Time by triangle:", "{:.5e}".format(t/ntri),"second.")
-f=open("RunningOn"+socket.gethostname(),"w")   
-f.write(str(t/ntri)+"\n")
-f.close()
-print("fin")
+ntri = 1_000_000
+random_generator = Rando()
+
+times = np.empty(ntri)
+
+for tr in range(ntri):
+    random_triangle(random_generator, x, y)
+    t1 = time()
+    stiffness.op(x, y, mat)
+    times[tr] = time() - t1
+
+if time.__name__.endswith("_ns"):
+    times *= 1e-9
+
+print(f"Time by triangle: {np.median(times) * 1e6:.2f} µs. (median)")
+print(f"Time by triangle: {np.min(times) * 1e6:.2f} µs. (min)")
+
+from transonic.util import timeit
+
+time_timeit = timeit("stiffness.op(x, y, mat)", globals=locals())
+
+print(f"Time by triangle: {time_timeit * 1e6:.2f} µs. (timeit)")
+
+with open("RunningOn" + socket.gethostname(), "w") as file:
+    file.write(f"{time_timeit}\n")

FeStiff/Pythran/rando.py

@@ -1,15 +1,19 @@
-class rando:
+class Rando:
     def __init__(self):
-        self.seed=123456789
-        self.a=1103515245
-        self.c=12345
-        self.m=2**32
+        self.seed = 123456789
+        self.a = 1103515245
+        self.c = 12345
+        self.m = 2 ** 32
+
     def get(self):
-        self.seed= (self.a * self.seed + self.c) % self.m
+        self.seed = (self.a * self.seed + self.c) % self.m
         return self.seed
-    def fv(self,vmax=1.):
-        return vmax*float(self.get())/self.m
+
+    def fv(self, vmax=1.0):
+        return vmax * float(self.get()) / self.m
+
+
 if __name__ == "__main__":
-    R=rando()
-    for i in range(0,100):
-        print(R.fv(10.))
+    R = rando()
+    for i in range(100):
+        print(R.fv(10.0))

FeStiff/Pythran/random_triangle.py

@@ -0,0 +1,6 @@
+def random_triangle(R, x, y):
+    # domain is [0,10]x[0,10]
+    for i in range(3):
+        x[i] = R.fv(10.0)
+    for i in range(3):
+        y[i] = R.fv(10.0)

FeStiff/Pythran/script

@@ -1,4 +1,4 @@
 #!/bin/bash
-pythran -march=native -O3 -DNDEBUG  StiffOut.py
+pythran -march=native -O3 -DUSE_XSIMD stiff_out.py
 echo "run test:"
 python3 main.py

FeStiff/Pythran/stiff_out.py

@@ -0,0 +1,26 @@
+
+# pythran export stiff_out(float[:],float[:],float[:],float[:, :, :],float[:, :, :])
+def stiff_out(x, y, out, grads, gq):
+    a11 = -y[0] + y[2]
+    a12 = y[0] - y[1]
+    a21 = x[0] - x[2]
+    a22 = -x[0] + x[1]
+
+    for f in range(6):
+        for p in range(3):
+            grads[f, p, 0] = a11 * gq[f, p, 0] + a12 * gq[f, p, 1]
+            grads[f, p, 1] = a21 * gq[f, p, 0] + a22 * gq[f, p, 1]
+
+    det = -(x[1] - x[2]) * y[0] + (x[0] - x[2]) * y[1] - (x[0] - x[1]) * y[2]
+    dv = 1.0 / (6.0 * det)
+    ii = 0
+    for i in range(6):
+        for j in range(i + 1):
+            s = 0.0
+            for k in range(3):
+                s += (
+                    grads[i, k, 0] * grads[j, k, 0]
+                    + grads[i, k, 1] * grads[j, k, 1]
+                )
+            out[ii] = dv * s
+            ii += 1

FeStiff/Pythran/stiffness.py

@@ -0,0 +1,24 @@
+import numpy as np
+
+from stiff_out import stiff_out
+
+
+class Stiffness:
+    __slots__ = ["grads", "gq"]
+
+    def __init__(self):
+        self.grads = np.empty((6, 3, 2))
+
+        # fmt: off
+        gq = np.array(
+            [-1.0, -1.0, 1.0, 1.0, -1.0, -1.0, 1.0, 0.0, 1.0, 0.0, -1.0,
+            0.0, 0.0, -1.0, 0.0, 1.0, 0.0, 1.0, 0.0, -2.0, -2.0, -2.0,
+            2.0, 0.0, 0.0, 2.0, 2.0, 2.0, 2.0, 0.0, 0.0, 2.0, -2.0, -2.0,
+            -2.0, 0.0]
+        )
+        # fmt: on
+
+        self.gq = gq.reshape(self.grads.shape)
+
+    def op(self, x, y, m):
+        stiff_out(x, y, m, self.grads, self.gq)