adding DecompCheck.py

DecompCheck.py

+import numpy as N
+import sympy.ntheory as SN
+import itertools
+
+####################
+# LOW LEVEL ROUTINES
+####################
+
+# independent of dimension
+
+# take out trivial divisors 1 and n
+def divisors(n):
+    return SN.divisors(n)[1:-1]
+
+# non-trivial decompositions of integer "n" into "dim" integer factors
+# returns an arbitrary length list of lists of length "dim" 
+# possibly with repeats
+def decompList(n, dim):
+    dl = divisors(n)
+    if dim == 2:
+        ret = []
+        for d in dl:
+            tmp = [d, n/d]
+            tmp.sort()
+            ret.append(tmp)
+        ret.sort()
+    else:
+        ret = []
+        for d in dl:
+            tmp = decompList(n/d, dim-1)
+            for i in tmp:
+                i.append(d)
+                i.sort()
+                ret.append(i)
+    return ret
+
+# returns a list of a unique set of rank-"dim" tuples
+def decompSet(n, dim, reverse=False):
+    ret = decompList(n,dim)
+    ret = map(tuple, ret)
+    ret = set(ret)
+    ret = list(ret)
+    ret.sort(reverse=reverse)
+    return ret
+
+# decompositions sorted by increasing aspect ratio of largest/smallest factor
+# can also limit by maximum aspect ratio
+def decompAspectRatio(n, dim, maxAspectRatio=0.0):
+    ret = decompSet(n, dim, False)
+    ret = map(lambda x: [1.0*x[-1]/x[0],x], ret)
+    ret.sort()
+    if maxAspectRatio > 0.0:
+        ret = filter(lambda x: x[0] <= maxAspectRatio, ret)
+    return ret
+
+# decompositions based on ranks, independent of grid dimensions
+
+# check whether a 2D decomp is commensurate with subset of 3D decomp
+# checks both orders of 2D decomp and returns True if either works
+# a = 3D subset, 2-tuple
+# b = 2D 2-tuple
+def check2d2d(a, b):
+    return (b[0]%a[0]==0 and b[1]%a[1]==0) or (b[1]%a[0]==0 and b[0]%a[1]==0)
+
+# filter list of 2D decomps for those commensurate with subset of 3D decomp
+# a = 3D subset 2-tuple
+# b = list of 2D 2-tuples
+def filter2d2d(a, blist):
+    return filter(lambda x: check2d2d(a,x), blist)
+
+# return xpencils, ypencils, and zpencils commensurate with a 3D decomp
+# pencil lists could be emtpy if there are no co-commensurate decomps
+# a = 3D 3-tuple
+def filter3d(a):
+    nranks = a[0]*a[1]*a[2]
+    blist = decompSet(nranks, 2)
+    blist.sort(reverse=True)
+    xpencils = filter2d2d((a[1],a[2]), blist)
+    ypencils = filter2d2d((a[2],a[0]), blist)
+    zpencils = filter2d2d((a[0],a[1]), blist)
+    return [xpencils, ypencils, zpencils]
+
+
+
+# sum the remainders of ng%tuple
+# works for arbitrary-rank tuples
+def sumRemainders(ng, tuple):
+    return N.array(map(lambda x: ng%x, tuple)).sum()
+
+# find pencil decompositions co-commensurate with 3D decomp and ng
+def filterNg3d2d(ng, tuple3d, filter3doutput):
+    if sumRemainders(ng, tuple3d) == 0:
+        xpencils = filter(lambda x: sumRemainders(ng, x)==0, filter3doutput[0])
+        ypencils = filter(lambda x: sumRemainders(ng, x)==0, filter3doutput[1])
+        zpencils = filter(lambda x: sumRemainders(ng, x)==0, filter3doutput[2])
+        if len(xpencils) > 0 and len(ypencils) > 0 and len(zpencils) > 0:
+            return [xpencils, ypencils, zpencils]
+    return []
+
+######################
+# CONVENIENCE ROUTINES
+######################
+
+def filterNg3d(ng, tuple3d):
+    filter3doutput = filter3d(tuple3d)
+    return filterNg3d2d(ng, tuple3d, filter3doutput)
+
+# find ng in range [ng0,ng1] that have pencil decomps for a given 3D decomp
+def ngCandidates3d(tuple3d, ng0, ng1):
+    ret = []
+    filter3doutput = filter3d(tuple3d)
+    for i in xrange(ng0, ng1+1):
+        if len(filterNg3d2d(i, tuple3d, filter3doutput)) > 0:
+            ret.append(i)
+    return ret
+
+# cannot remember what permutations does, seems unused
+def permutations(tuple):
+    ret = list(set(list(itertools.permutations(tuple))))
+    ret.sort()
+    return ret
+
+# cannot remember what doublecheck does, seems unused
+def doublecheck(nranks, ng, t3d, t2dx, t2dy, t2dz):
+    # decomps have right number of ranks
+    if N.array(t3d).prod() != nranks:
+        return False
+    if N.array(t2dx).prod() != nranks:
+        return False
+    if N.array(t2dy).prod() != nranks:
+        return False
+    if N.array(t2dz).prod() != nranks:
+        return False
+    # 2D decomps are commensurate with 3D decomps
+    if t2dx[1] % t3d[1] != 0 or t2dx[2] % t3d[2] != 0:
+        return False
+    if t2dy[0] % t3d[0] != 0 or t2dy[2] % t3d[2] != 0:
+        return False
+    if t2dz[0] % t3d[0] != 0 or t2dz[1] % t3d[1] != 0:
+        return False
+    # decomps are divisors of ng
+    if sumRemainders(ng, t3d) != 0:
+        return False
+    if sumRemainders(ng, t2dx) != 0:
+        return False
+    if sumRemainders(ng, t2dy) != 0:
+        return False
+    if sumRemainders(ng, t2dz) != 0:
+        return False
+    # ok looks fine
+    return True
+
+#####################
+# HIGH LEVEL ROUTINES
+#####################
+
+# find ng in range [ng0,ng1] that have pencil decomps for any valid 3D decomp
+# optionally specify maximum 3D decomp aspect ratio to check
+def ngCandidatesNranks(nranks, ng0, ng1, maxAspectRatio=0.0):
+    tuple3dlist = decompAspectRatio(nranks, 3, maxAspectRatio)
+    ret = []
+    for i in tuple3dlist:
+        tuple3d = i[1]
+        ret += ngCandidates3d(tuple3d, ng0, ng1)
+    ret = set(ret)
+    ret = list(ret)
+    ret.sort()
+    return ret
+
+# find 3D decomps that have co-commensurate pencil decomps for ng and nranks
+# optionally specify maximum 3D decomp aspect ratio to check
+def filterNgNranks(ng, nranks, maxAspectRatio=0.0):
+    list3d = decompAspectRatio(nranks, 3, maxAspectRatio)
+    return filter(lambda x: len(filterNg3d(ng, x[1])) > 0, list3d)
+
+# expand pencils into 3D tuples in all orders that work
+def enumeratePencilsNg3d(ng, tuple3d):
+    pencils = filterNg3d(ng, tuple3d)
+    xret = []
+    for pencil in pencils[0]:
+        if pencil[0]%tuple3d[1]==0 and pencil[1]%tuple3d[2]==0:
+            xret.append( (1, pencil[0], pencil[1]) )
+        if pencil[1]%tuple3d[1]==0 and pencil[0]%tuple3d[2]==0:
+            xret.append( (1, pencil[1], pencil[0]) )
+    yret = []
+    for pencil in pencils[1]:
+        if pencil[0]%tuple3d[0]==0 and pencil[1]%tuple3d[2]==0:
+            yret.append( (pencil[0], 1, pencil[1]) )
+        if pencil[1]%tuple3d[0]==0 and pencil[0]%tuple3d[2]==0:
+            yret.append( (pencil[1], 1, pencil[0]) )
+    zret = []
+    for pencil in pencils[2]:
+        if pencil[0]%tuple3d[0]==0 and pencil[1]%tuple3d[1]==0:
+            zret.append( (pencil[0], pencil[1], 1) )
+        if pencil[1]%tuple3d[0]==0 and pencil[0]%tuple3d[1]==0:
+            zret.append( (pencil[1], pencil[0], 1) )
+    return [xret, yret, zret]