bench: cleanup redundant code and add Operator and XDSLOperator at same file

devito/ir/ietxdsl/cluster_to_ssa.py

@@ -113,7 +113,7 @@ def _convert_eq(self, eq: LoweredEq):
         ), f"can only write to offset [0,0,0], given {offsets[1:]}"
 
         self.block.add_op(stencil.ReturnOp.get([rhs_result]))
-        outermost_block.add_op(func.Return.get())
+        outermost_block.add_op(func.Return())
 
         return func.FuncOp.from_region(
             "apply_kernel", [], [], Region([outermost_block])

devito/ir/ietxdsl/xdsl_passes.py

@@ -143,7 +143,7 @@ def _op_to_func(op: Operator):
             ietxdsl_functions.myVisit(i, block=block, ssa_vals=ssa_val_dict)
 
     # add a trailing return
-    block.add_op(func.Return.get())
+    block.add_op(func.Return())
 
     func_op = func.FuncOp.from_region(str(op.name), arg_types, [], Region([block]))
 

devito/operator/xdsl_operator.py

@@ -46,12 +46,12 @@
 """
 
 
-CFLAGS = "-O3 -march=native -mtune=native"
+CFLAGS = "-O3 -march=native -mtune=native -lmlir_c_runner_utils"
 
-MLIR_CPU_PIPELINE = '"builtin.module(canonicalize, cse, loop-invariant-code-motion, canonicalize, cse, loop-invariant-code-motion,cse,canonicalize,fold-memref-alias-ops,lower-affine,convert-scf-to-cf,convert-math-to-llvm,convert-func-to-llvm{use-bare-ptr-memref-call-conv},finalize-memref-to-llvm,canonicalize,cse)"'
-MLIR_OPENMP_PIPELINE = '"builtin.module(canonicalize, cse, loop-invariant-code-motion, canonicalize, cse, loop-invariant-code-motion,cse,canonicalize,fold-memref-alias-ops,lower-affine,finalize-memref-to-llvm,loop-invariant-code-motion,canonicalize,cse,convert-scf-to-openmp,finalize-memref-to-llvm,convert-scf-to-cf,convert-func-to-llvm{use-bare-ptr-memref-call-conv},convert-openmp-to-llvm,convert-math-to-llvm,reconcile-unrealized-casts,canonicalize,cse)"'
+MLIR_CPU_PIPELINE = '"builtin.module(canonicalize, cse, loop-invariant-code-motion, canonicalize, cse, loop-invariant-code-motion,cse,canonicalize,fold-memref-alias-ops,expand-strided-metadata, loop-invariant-code-motion,lower-affine,convert-scf-to-cf,convert-math-to-llvm,convert-func-to-llvm{use-bare-ptr-memref-call-conv},finalize-memref-to-llvm,canonicalize,cse)"'
+MLIR_OPENMP_PIPELINE = '"builtin.module(canonicalize, cse, loop-invariant-code-motion, canonicalize, cse, loop-invariant-code-motion,cse,canonicalize,fold-memref-alias-ops,expand-strided-metadata, loop-invariant-code-motion,lower-affine,finalize-memref-to-llvm,loop-invariant-code-motion,canonicalize,cse,convert-scf-to-openmp,finalize-memref-to-llvm,convert-scf-to-cf,convert-func-to-llvm{use-bare-ptr-memref-call-conv},convert-openmp-to-llvm,convert-math-to-llvm,reconcile-unrealized-casts,canonicalize,cse)"'
 # gpu-launch-sink-index-computations seemed to have no impact
-MLIR_GPU_PIPELINE = '"builtin.module(test-math-algebraic-simplification,scf-parallel-loop-tiling{parallel-loop-tile-sizes=128,1,1},func.func(gpu-map-parallel-loops),convert-parallel-loops-to-gpu,fold-memref-alias-ops,lower-affine,gpu-kernel-outlining,canonicalize,cse,convert-arith-to-llvm{index-bitwidth=64},finalize-memref-to-llvm{index-bitwidth=64},convert-scf-to-cf,convert-cf-to-llvm{index-bitwidth=64},canonicalize,cse,gpu.module(convert-gpu-to-nvvm,reconcile-unrealized-casts,canonicalize,gpu-to-cubin),gpu-to-llvm,canonicalize,cse)"'
+MLIR_GPU_PIPELINE = '"builtin.module(test-math-algebraic-simplification,scf-parallel-loop-tiling{parallel-loop-tile-sizes=128,1,1},func.func(gpu-map-parallel-loops),convert-parallel-loops-to-gpu,fold-memref-alias-ops,expand-strided-metadata,lower-affine,gpu-kernel-outlining,canonicalize,cse,convert-arith-to-llvm{index-bitwidth=64},finalize-memref-to-llvm{index-bitwidth=64},convert-scf-to-cf,convert-cf-to-llvm{index-bitwidth=64},canonicalize,cse,gpu.module(convert-gpu-to-nvvm,reconcile-unrealized-casts,canonicalize,gpu-to-cubin),gpu-to-llvm,canonicalize,cse)"'
 
 XDSL_CPU_PIPELINE = "stencil-shape-inference,convert-stencil-to-ll-mlir,printf-to-llvm"
 XDSL_GPU_PIPELINE = "stencil-shape-inference,convert-stencil-to-ll-mlir{target=gpu},printf-to-llvm"
@@ -137,7 +137,10 @@ def _jit_compile(self):
                 print("JIT Backdoor: loading xdsl file from: " + backdoor)
                 with open(backdoor, 'r') as f:
                     module_str = f.read()
-
+            source_name = os.path.splitext(self._tf.name)[0] + ".mlir"
+            source_file = open(source_name, "w")
+            source_file.write(module_str)
+            source_file.close()
             # compile IR using xdsl-opt | mlir-opt | mlir-translate | clang
             try:
                 cflags = CFLAGS
@@ -151,17 +154,20 @@ def _jit_compile(self):
                 if is_gpu:
                     cflags += " -lmlir_cuda_runtime "
 
-                cmd = f'xdsl-opt -p {xdsl_pipeline} |' \
+                # TODO More detailed error handling manually,
+                # instead of relying on a bash-only feature.
+                cmd = 'set -eo pipefail; '\
+                    f'xdsl-opt {source_name} -p {xdsl_pipeline} |' \
                     f'mlir-opt -p {mlir_pipeline} | ' \
                     f'mlir-translate --mlir-to-llvmir | ' \
                     f'{cc} {cflags} -shared -o {self._tf.name} {self._interop_tf.name} -xir -'
-
+                print(cmd)
                 res = subprocess.run(
                     cmd,
                     shell=True,
-                    input=module_str,
                     text=True,
                     capture_output=True,
+                    executable="/bin/bash"
                 )
 
                 if res.returncode != 0:
@@ -177,7 +183,7 @@ def _jit_compile(self):
         elapsed = self._profiler.py_timers['jit-compile']
 
         perf("XDSLOperator `%s` jit-compiled `%s` in %.2f s with `mlir-opt`" %
-                    (self.name, self._tf.name, elapsed))
+                    (self.name, source_name, elapsed))
 
     @property
     def _soname(self):

fast/diffusion_2D_wBCs.py

@@ -42,49 +42,28 @@
 u = TimeFunction(name='u', grid=grid, space_order=so)
 
 # Reset our data field and ICs
-# u.data[:, :, :] = 0.1
 init_hat(field=u.data[0], dx=dx, dy=dy, value=1.)
 
-
-# u.data[0, :, int(ny/2)] = 2
-
 a = Constant(name='a')
 # Create an equation with second-order derivatives
-eq = Eq(u.dt, a * u.laplace, subdomain=grid.interior)
+# eq = Eq(u.dt, a * u.laplace, subdomain=grid.interior)
+eq = Eq(u.dt, a * u.laplace)
 stencil = solve(eq, u.forward)
 eq_stencil = Eq(u.forward, stencil)
 
 # Create boundary condition expressions
 x, y = grid.dimensions
 t = grid.stepping_dim
 
-# Add boundary conditions
-# bc = [Eq(u[t+1, x, y, 0], 2.)]  # bottom
-# bc += [Eq(u[t+1, x, y, nz-1], 2.)]  # top
-# bc += [Eq(u[t+1, 0, y, z], 2.)]  # left
-# bc += [Eq(u[t+1, nx-1, y, z], 2.)]  # right
-
-# bc += [Eq(u[t+1, x, 0, z], 2.)]  # front
-# bc += [Eq(u[t+1, x, ny-1, z], 2.)]  # back
-
-print(eq_stencil)
-
-# Create an operator that updates the forward stencil point
-# plus adding boundary conditions
-# op = Operator([eq_stencil] + bc, subdomain=grid.interior)
-
-# No BCs
-# op = XDSLOperator([eq_stencil])
-op = Operator([eq_stencil])
-# print(op.ccode)
-# dt = 0.00002
-# import pdb;pdb.set_trace()
-# Apply the operator for a number of timesteps
-print(dt)
+initdata = u.data[:]
+op = Operator([eq_stencil], name='DevitoOperator')
 op.apply(time=nt, dt=dt, a=nu)
 
-print("Field norm is:", norm(u))
+print("Devito Field norm is:", norm(u))
 
-plot_image(u.data[0], cmap="seismic")
+# Reset data and run XDSLOperator
+init_hat(field=u.data[0], dx=dx, dy=dy, value=1.)
+xdslop = XDSLOperator([eq_stencil], name='XDSLOperator')
+xdslop.apply(time=nt, dt=dt, a=nu)
 
-# import pdb;pdb.set_trace()
+print("XDSL Field norm is:", norm(u))

fast/diffusion_3D_wBCs.py

@@ -65,39 +65,33 @@ def plot_3dfunc(u):
 
 a = Constant(name='a')
 # Create an equation with second-order derivatives
-eq = Eq(u.dt, a * u.laplace, subdomain=grid.interior)
+eq = Eq(u.dt, a * u.laplace)
+
 stencil = solve(eq, u.forward)
 eq_stencil = Eq(u.forward, stencil)
 
 # Create boundary condition expressions
 x, y, z = grid.dimensions
 t = grid.stepping_dim
 
-# Add boundary conditions
-# bc = [Eq(u[t+1, x, y, 0], 2.)]  # bottom
-# bc += [Eq(u[t+1, x, y, nz-1], 2.)]  # top
-# bc += [Eq(u[t+1, 0, y, z], 2.)]  # left
-# bc += [Eq(u[t+1, nx-1, y, z], 2.)]  # right
-
-# bc += [Eq(u[t+1, x, 0, z], 2.)]  # front
-# bc += [Eq(u[t+1, x, ny-1, z], 2.)]  # back
-
 print(eq_stencil)
 
-# Create an operator that updates the forward stencil point
-# plus adding boundary conditions
-# op = Operator([eq_stencil] + bc, subdomain=grid.interior)
-
-# No BCs
-# op = XDSLOperator([eq_stencil])
-op = Operator([eq_stencil])
-# print(op.ccode)
-
+# Create Operator
+op = Operator([eq_stencil], name='DevitoOperator')
 # Apply the operator for a number of timesteps
 op.apply(time=nt, dt=dt, a=nu)
+print("Devito Field norm is:", norm(u))
+
+# Reset field
+u.data[:, :, :, :] = 0
+u.data[:, :, :, int(nz/2)] = 1
+xdslop = XDSLOperator([eq_stencil], name='xDSLOperator')
+# Apply the xdsl operator for a number of timesteps
+xdslop.apply(time=nt, dt=dt, a=nu)
 
 if args.plot:
     plot_3dfunc(u)
 
-print("Field norm is:", norm(u))
+print("XDSL Field norm is:", norm(u))
+
 # import pdb;pdb.set_trace()

fast/nd_nwave_devito_nodamp.py

@@ -102,30 +102,30 @@ def plot_3dfunc(u):
 
 # Finally we define the source injection and receiver read function to generate
 # the corresponding code
-src_term = src.inject(field=u.forward, expr=src * dt**2 / model.m)
-op = Operator([stencil] + src_term, subs=model.spacing_map)
-
 print(time_range)
+src_term = src.inject(field=u.forward, expr=src * dt**2 / model.m)
+op = Operator([stencil] + src_term, subs=model.spacing_map, name='DevitoOperator')
 # Run with source and plot
-# import pdb;pdb.set_trace()
 op.apply(time=time_range.num-1, dt=model.critical_dt)
 
 if len(shape) == 3:
     if args.plot:
         plot_3dfunc(u)
 
-# Run more with no sources now (Not supported in xdsl)
-# op = XDSLOperator([stencil])
+initdata = u.data[:]
 
-# op.apply(time=time_range.num-1, dt=model.critical_dt)
+# Run more with no sources now (Not supported in xdsl)
+xdslop = XDSLOperator([stencil], name='XDSLOperator')
+xdslop.apply(time=time_range.num-1, dt=model.critical_dt)
 
 if len(shape) == 3:
     if args.plot:
         plot_3dfunc(u)
 
-# print(norm(u))
+print(norm(u))
 
-# todo edit data
+# Reset initial data
+u.data[:] = initdata
 
 # Run more with no sources now (Not supported in xdsl)
 xdslop = XDSLOperator([stencil])