Add proper flag handling.

xdslproject · Jul 31, 2023 · 90cb3fd · 90cb3fd
1 parent b135fba
commit 90cb3fd
Showing 1 changed file with 18 additions and 16 deletions.
diff --git a/devito/ir/ietxdsl/cluster_to_ssa.py b/devito/ir/ietxdsl/cluster_to_ssa.py
@@ -1,6 +1,8 @@
 # ------------- devito import -------------#
 
 from sympy import Add, Expr, Float, Indexed, Integer, Mod, Mul, Pow, Symbol
+from devito.arch.archinfo import NvidiaDevice
+from devito.parameters import configuration
 from xdsl.dialects import arith, builtin, func, memref, scf, stencil, gpu
 from xdsl.dialects.experimental import dmp, math
 from xdsl.ir import Attribute, Block, Operation, OpResult, Region, SSAValue
@@ -480,22 +482,22 @@ def match_and_rewrite(self, op: iet_ssa.LoadSymbolic, rewriter: PatternRewriter,
             body = parent.body.blocks[0]
             args[symb_name] = body.insert_arg(op.result.type, len(body.args))
             # GPU STUFF
-
-            arg = args[symb_name]
-            shapetype = arg.type
-            if isinstance(shapetype, stencil.FieldType):
-                memref_type = memref.MemRefType.from_element_type_and_shape(shapetype.get_element_type(), shapetype.get_shape())
-                alloc = gpu.AllocOp(memref.MemRefType.from_element_type_and_shape(shapetype.get_element_type(), shapetype.get_shape()))
-                outcast = builtin.UnrealizedConversionCastOp.get(alloc, shapetype)
-                # arg.replace_by(outcast.results[0])
-                args[symb_name] = outcast.results[0]
-                incast = builtin.UnrealizedConversionCastOp.get(arg, memref_type)
-                copy = gpu.MemcpyOp(source=incast, destination=alloc)
-                body.insert_ops_before([alloc, outcast, incast, copy], body.ops.first)
-
-                copy_out = gpu.MemcpyOp(source=alloc, destination=incast)
-                dealloc = gpu.DeallocOp(alloc)
-                body.insert_ops_before([copy_out, dealloc], body.ops.last)
+            if isinstance(configuration['platform'], NvidiaDevice):
+                arg = args[symb_name]
+                shapetype = arg.type
+                if isinstance(shapetype, stencil.FieldType):
+                    memref_type = memref.MemRefType.from_element_type_and_shape(shapetype.get_element_type(), shapetype.get_shape())
+                    alloc = gpu.AllocOp(memref.MemRefType.from_element_type_and_shape(shapetype.get_element_type(), shapetype.get_shape()))
+                    outcast = builtin.UnrealizedConversionCastOp.get(alloc, shapetype)
+                    # arg.replace_by(outcast.results[0])
+                    args[symb_name] = outcast.results[0]
+                    incast = builtin.UnrealizedConversionCastOp.get(arg, memref_type)
+                    copy = gpu.MemcpyOp(source=incast, destination=alloc)
+                    body.insert_ops_before([alloc, outcast, incast, copy], body.ops.first)
+
+                    copy_out = gpu.MemcpyOp(source=alloc, destination=incast)
+                    dealloc = gpu.DeallocOp(alloc)
+                    body.insert_ops_before([copy_out, dealloc], body.ops.last)
 
         op.result.replace_by(args[symb_name])