compiler: Tweak linearize pass

Author: FabioLuporini

devito/passes/iet/linearization.py

@@ -223,7 +223,7 @@ def linearize_accesses(iet, key0, tracker=None):
 
     # 2) What `iet` *offers*
     # E.g. `{x_fsz0 -> u_vec->size[1]}`
-    defines = FindSymbols('defines-aliases').visit(iet)
+    defines = FindSymbols('defines').visit(iet)
     offers = filter_ordered(i for i in defines if key0(i.function))
     instances = {}
     for i in offers:
@@ -294,16 +294,9 @@ def _(f, d):
 
 
 @_generate_fsz.register(Array)
-def _(f, d):
-    return f.symbolic_shape[d]
-
-
 @_generate_fsz.register(Bundle)
 def _(f, d):
-    if f.is_DiscreteFunction:
-        return _generate_fsz.registry[DiscreteFunction](f, d)
-    else:
-        return _generate_fsz.registry[Array](f, d)
+    return f.symbolic_shape[d]
 
 
 @singledispatch

tests/test_linearize.py

@@ -342,20 +342,21 @@ def test_strides_forwarding1():
     linearize(graph, callback=True, options={'index-mode': 'int32'},
               sregistry=SymbolRegistry())
 
-    # Despite `a` is passed via `a.indexed`, and since it's an Array (which
-    # have symbolic shape), we expect the stride exprs to be placed in `bar`,
-    # and in `bar` only, as `foo` doesn't really use `a`, it just propagates it
-    # down to `bar`
+    # `a` is passed via `a.indexed`, so the stride exprs are expected to be
+    # placed in `foo` and then passed down to `bar` as arguments
     foo = graph.root
     bar = graph.efuncs['bar']
 
     assert len(foo.body.body) == 1
     assert foo.body.body[0].is_Call
+    assert len(foo.body.strides) == 3
+    assert foo.body.strides[0].write.name == 'y_fsz0'
+    assert foo.body.strides[2].write.name == 'y_stride0'
 
+    assert len(bar.parameters) == 2
+    assert bar.parameters[0].name == 'a'
+    assert bar.parameters[1].name == 'y_stride0'
     assert len(bar.body.body) == 1
-    assert len(bar.body.strides) == 3
-    assert bar.body.strides[0].write.name == 'y_fsz0'
-    assert bar.body.strides[2].write.name == 'y_stride0'
 
 
 def test_strides_forwarding2():