Skip to content

Commit 5208924

Browse files
ZoeLeibowitzZoeLeibowitz
committed
need to fix with space dimensions
1 parent 8a09272 commit 5208924

8 files changed

Lines changed: 2211 additions & 2037 deletions

File tree

‎devito/ir/iet/nodes.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1148,8 +1148,8 @@ def expr_symbols(self):
11481148
ret.extend([self.pointer._C_symbol, self.pointee._C_symbol])
11491149
else:
11501150
ret.extend([self.pointer, self.pointee.indexed])
1151-
ret.extend(flatten(i.free_symbols
1152-
for i in self.pointee.symbolic_shape[1:]))
1151+
ret.extend(flatten(i.free_symbols
1152+
for i in self.pointee.symbolic_shape[1:]))
11531153
else:
11541154
assert False, f"Unexpected pointer type {type(self.pointer)}"
11551155

‎devito/petsc/equations.py‎

Lines changed: 139 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from sympy import Eq
22
from devito.symbolics import retrieve_indexed, retrieve_dimensions
33
from devito.petsc.types.equation import ConstrainBC
4-
from devito.types.dimension import CustomBoundSubDimension
4+
from devito.types.dimension import CustomBoundSubDimension, SpaceDimension, CustomBoundSpaceDimension
55
from devito import Min, Max
66

77

@@ -26,13 +26,17 @@ def constrain_essential_bcs(expressions, **kwargs):
2626
new_exprs.append(e)
2727
continue
2828

29+
30+
# this needs to be applied to all space dimensions and all subdims that are "MIDDLE"
31+
2932
indexeds = retrieve_indexed(e)
3033
dims = retrieve_dimensions([i for j in indexeds for i in j.indices], mode='unique')
3134
# implicit_dims = set(e.implicit_dims)
3235
dims.update(e.implicit_dims)
3336
# from IPython import embed; embed()
3437
dims = [d for d in dims if d.is_Sub and not d.local]
3538

39+
3640
if not dims:
3741
new_exprs.append(e)
3842
continue
@@ -80,4 +84,137 @@ def constrain_essential_bcs(expressions, **kwargs):
8084
new_e = new_e._rebuild(implicit_dims=implicit_dims_new)
8185
new_exprs.append(new_e)
8286

83-
return new_exprs
87+
return new_exprs
88+
89+
90+
91+
# def constrain_essential_bcs(expressions, **kwargs):
92+
# """TODO: improve docs ..Modify the subdims used in ConstrainEssentialBC equations ... to locally
93+
# constrain nodes (including non owned halo nodes) ....."""
94+
95+
# mapper = {}
96+
# new_exprs = []
97+
98+
# # build mapper
99+
# for e in expressions:
100+
# if not isinstance(e, ConstrainBC):
101+
# # new_exprs.append(e)
102+
# continue
103+
# # this needs to be applied to all space dimensions and all subdims that are "MIDDLE"
104+
105+
# indexeds = retrieve_indexed(e)
106+
# dims = retrieve_dimensions([i for j in indexeds for i in j.indices], mode='unique')
107+
# # from IPython import embed; embed()
108+
# # implicit_dims = set(e.implicit_dims)
109+
# dims.update(e.implicit_dims)
110+
# # from IPython import embed; embed()
111+
112+
# # Collect non-local subdims (which would be generated from a "middle" subdomain)
113+
# # "local" subdims are generated from 'left' and 'right" subdomains and do not need to be considered
114+
# # because by construction, left and right subdimensions cannot cross ranks
115+
# subdims = [d for d in dims if d.is_Sub and not d.local]
116+
117+
# space_dims = [d for d in dims if isinstance(d, SpaceDimension)]
118+
119+
# relevant_dims = subdims + space_dims
120+
# # relevant_dims = subdims
121+
122+
123+
# if not relevant_dims:
124+
# continue
125+
126+
127+
# # first map the subdims
128+
# for d in subdims:
129+
# # replace the dim with a new one that has a different symbolic_min and symbolic_max
130+
131+
# # obvs shouldn't be obtained from indexeds[0]
132+
# # USE e.lhs function -> the one that the BC is being applied to
133+
# # from IPython import embed; embed()
134+
# # f._size_nodomain.left
135+
# # from IPython import embed; embed()
136+
# # halo_size_left = indexeds[0].function._size_halo[d].left
137+
# # halo_size_right = indexeds[0].function._size_halo[d].right
138+
139+
# halo_size_left = 2
140+
# halo_size_right = 2
141+
142+
# from devito.petsc.types.dimension import SubDimMax, SubDimMin
143+
144+
# # TODO: change name..
145+
146+
# # in theory this class shoulod just take in d
147+
# # TODO: use unique name
148+
# sregistry = kwargs.get('sregistry')
149+
# subdim_max = SubDimMax(sregistry.make_name(prefix=d.name + '_max'), subdim=d)
150+
# subdim_min = SubDimMin(sregistry.make_name(prefix=d.name + '_min'), subdim=d)
151+
# # from IPython import embed; embed()
152+
# # unique_name
153+
# new_dim = CustomBoundSubDimension(
154+
# name=d.name,
155+
# parent=d.parent,
156+
# thickness=d.thickness,
157+
# local=d.local,
158+
# custom_left=Max(subdim_min, d.parent.symbolic_min - halo_size_left),
159+
# custom_right=Min(subdim_max, d.parent.symbolic_max + halo_size_right)
160+
# )
161+
# mapper[d] = new_dim
162+
163+
# # then tackle space dims
164+
# for d in space_dims:
165+
# mapper[d] = d
166+
# # for d in space_dims:
167+
# # halo_size_left = 2
168+
# # halo_size_right = 2
169+
170+
# # from devito.petsc.types.dimension import SpaceDimMax, SpaceDimMin
171+
172+
# # # TODO: change name..
173+
174+
# # # in theory this class shoulod just take in d
175+
# # # TODO: use unique name
176+
# # sregistry = kwargs.get('sregistry')
177+
# # space_dim_max = SpaceDimMax(sregistry.make_name(prefix=d.name + '_max'), space_dim=d)
178+
# # space_dim_min = SpaceDimMin(sregistry.make_name(prefix=d.name + '_min'), space_dim=d)
179+
180+
# # # unique_name
181+
# # new_dim = CustomBoundSpaceDimension(
182+
# # name=d.name,
183+
# # custom_left=Max(space_dim_min, d.symbolic_min - halo_size_left),
184+
# # custom_right=Min(space_dim_max, d.symbolic_max + halo_size_right)
185+
# # )
186+
# # mapper[d] = new_dim
187+
# # # from IPython import embed; embed()
188+
189+
190+
# # from IPython import embed; embed()
191+
192+
193+
# for e in expressions:
194+
195+
# if not isinstance(e, ConstrainBC):
196+
# new_exprs.append(e)
197+
# continue
198+
199+
# indexeds = retrieve_indexed(e)
200+
# dims = retrieve_dimensions([i for j in indexeds for i in j.indices], mode='unique')
201+
# # implicit_dims = set(e.implicit_dims)
202+
# dims.update(e.implicit_dims)
203+
# # from IPython import embed; embed()
204+
# subdims = [d for d in dims if d.is_Sub and not d.local]
205+
206+
# space_dims = [d for d in dims if isinstance(d, SpaceDimension)]
207+
208+
# relevant_dims = subdims + space_dims
209+
210+
# if not relevant_dims:
211+
# new_exprs.append(e)
212+
# continue
213+
214+
# new_e = e.subs(mapper)
215+
# if e.implicit_dims:
216+
# implicit_dims_new = tuple(mapper.get(d, d) for d in e.implicit_dims)
217+
# new_e = new_e._rebuild(implicit_dims=implicit_dims_new)
218+
# new_exprs.append(new_e)
219+
220+
# return new_exprs

‎devito/petsc/iet/builder.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -330,7 +330,7 @@ def _setup(self):
330330
snes_set_options,
331331
call_struct_callback,
332332
mat_set_dm,
333-
calls_set_app_ctx,
333+
# calls_set_app_ctx,
334334
create_field_decomp,
335335
matop_create_submats_op,
336336
call_coupled_struct_callback,
@@ -367,6 +367,9 @@ def _create_dmda_calls(self, dmda):
367367

368368
get_local_section = petsc_call('DMGetLocalSection', [dmda, Byref(sobjs['lsection'])])
369369

370+
import cgen as c
371+
tmp = c.Line("PetscCall(PetscSectionView(lsection0, NULL));")
372+
370373
get_point_sf = petsc_call('DMGetPointSF', [dmda, Byref(sobjs['sf'])])
371374

372375
create_global_section = petsc_call(
@@ -395,6 +398,7 @@ def _create_dmda_calls(self, dmda):
395398
count_bcs,
396399
set_point_bcs,
397400
get_local_section,
401+
tmp,
398402
get_point_sf,
399403
create_global_section,
400404
dm_set_global_section,

‎devito/petsc/iet/callbacks.py‎

Lines changed: 4 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -731,13 +731,6 @@ def _create_count_bc_body(self, body):
731731
# Replace non-function data with pointer to data in struct
732732
subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for
733733
i in fields if not isinstance(i.function, AbstractFunction)}
734-
735-
# subs[]
736-
# subs[self.target] = sobjs['numBC']
737-
738-
# subs[Counter._C_symbol] = Cast(Deref(sobjs['numBCPtr']._C_symbol))
739-
740-
# from IPython import embed; embed()
741734

742735
return Uxreplace(subs).visit(body)
743736

@@ -750,30 +743,10 @@ def _create_set_point_bc_body(self, body):
750743
dmda = sobjs['callbackdm']
751744
ctx = objs['dummyctx']
752745

753-
754746
dm_get_local_info = petsc_call(
755747
'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)]
756748
)
757749

758-
# import numpy as np
759-
# if self.options['index-mode'] == 'int32':
760-
# dtype = np.int32
761-
# else:
762-
# dtype = np.int64
763-
# from devito.passes.iet.linearization import Tracker
764-
765-
# tracker = Tracker('basic', dtype, self.sregistry)
766-
# # from IPython import embed; embed()
767-
# key = lambda f: f.name == 'u'
768-
# body = linearize_accesses(body, key0=key, tracker=tracker)
769-
770-
# # will only be findexeds 'indexeds'
771-
# findexeds = FindSymbols('indexeds').visit(body)
772-
# mapper_findexeds = {i: i.linear_index for i in findexeds}
773-
774-
# from IPython import embed; embed()
775-
776-
# findexeds =
777750
body = self.time_dependence.uxreplace_time(body)
778751

779752
fields = get_user_struct_fields(body)
@@ -784,8 +757,10 @@ def _create_set_point_bc_body(self, body):
784757
)
785758

786759
comm = sobjs['comm']
760+
# Obvs fix the first arg
787761
is_create_general = petsc_call(
788-
'ISCreateGeneral', [comm, sobjs['numBC'], sobjs['bcPointsArr'], 'PETSC_OWN_POINTER', Byref(sobjs['bcPointsIS'])]
762+
'ISCreateGeneral', ['PetscObjectComm((PetscObject)(dm0))', sobjs['numBC'], sobjs['bcPointsArr'],
763+
'PETSC_OWN_POINTER', Byref(sobjs['bcPointsIS'])]
789764
)
790765

791766
malloc_bc_points_arr = petsc_call(
@@ -824,10 +799,7 @@ def _create_set_point_bc_body(self, body):
824799

825800
subs[Counter._C_symbol] = sobjs['bcPointsArr'].indexed[sobjs['k_iter']]
826801

827-
# body = Uxreplace(mapper_findexeds).visit(body)
828-
body = Uxreplace(subs).visit(body)
829-
830-
return body
802+
return Uxreplace(subs).visit(body)
831803

832804
def _make_user_struct_callback(self):
833805
"""

‎devito/petsc/types/dimension.py‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,4 +63,74 @@ def _arg_values(self, grid=None, **kwargs):
6363
val = decomp.index_glb_to_loc_unsafe(g_x_m + gltkn)
6464

6565
return {self.name: int(val)}
66+
67+
68+
69+
class SpaceDimMax(Thickness):
70+
"""
71+
"""
72+
73+
def __init_finalize__(self, *args, **kwargs):
74+
self._space_dim = kwargs.pop('space_dim')
75+
self._dtype = self._space_dim.dtype
76+
77+
super().__init_finalize__(*args, **kwargs)
78+
79+
@property
80+
def space_dim(self):
81+
return self._space_dim
82+
83+
84+
def _arg_values(self, grid=None, **kwargs):
85+
86+
dist = grid.distributor
87+
88+
# global rtkn
89+
# grtkn = kwargs.get(self.subdim.rtkn.name, self.subdim.rtkn.value)
90+
# print(g_x_M)
91+
# decomposition info
92+
decomp = dist.decomposition[self.space_dim]
93+
# obvs not just x etc..
94+
g_x_M = decomp.glb_max
95+
# print(g_x_M)
96+
val = decomp.index_glb_to_loc_unsafe(g_x_M)
97+
# print(val)
98+
99+
100+
return {self.name: int(val)}
101+
102+
103+
104+
class SpaceDimMin(Thickness):
105+
"""
106+
"""
107+
108+
def __init_finalize__(self, *args, **kwargs):
109+
self._space_dim = kwargs.pop('space_dim')
110+
self._dtype = self._space_dim.dtype
111+
112+
super().__init_finalize__(*args, **kwargs)
113+
114+
@property
115+
def space_dim(self):
116+
return self._space_dim
117+
118+
119+
def _arg_values(self, grid=None, **kwargs):
120+
121+
dist = grid.distributor
122+
123+
124+
decomp = dist.decomposition[self.space_dim]
125+
# obvs not just x etc..
126+
g_x_m = decomp.glb_min
127+
# print(g_x_M)
128+
val = decomp.index_glb_to_loc_unsafe(g_x_m)
129+
# print(val)
130+
131+
132+
return {self.name: int(val)}
133+
134+
135+
66136

‎devito/types/dimension.py‎

Lines changed: 26 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -844,11 +844,32 @@ def custom_left(self):
844844
def custom_right(self):
845845
return self._custom_right
846846

847-
# @cached_property
848-
# def _interval(self):
849-
# left = self.parent.symbolic_min + self._offset_left
850-
# right = self.parent.symbolic_max - self._offset_right
851-
# return sympy.Interval(left, right)
847+
@cached_property
848+
def _interval(self):
849+
left = self.custom_left
850+
right = self.custom_right
851+
return sympy.Interval(left, right)
852+
853+
854+
class CustomBoundSpaceDimension(SpaceDimension):
855+
856+
# have is_CustomSub = True ... here?
857+
858+
__rargs__ = SpaceDimension.__rargs__ + ('custom_left', 'custom_right')
859+
860+
def __init_finalize__(self, name,
861+
custom_left=0, custom_right=0, **kwargs):
862+
self._custom_left = custom_left
863+
self._custom_right = custom_right
864+
super().__init_finalize__(name, **kwargs)
865+
866+
@property
867+
def custom_left(self):
868+
return self._custom_left
869+
870+
@property
871+
def custom_right(self):
872+
return self._custom_right
852873

853874
@cached_property
854875
def _interval(self):

0 commit comments

Comments
 (0)