[python][gpu] Add mem_fence function (#199)

Author: Hardcode84

dpcomp_gpu_runtime/lib/kernel_api_stubs.cpp

@@ -51,6 +51,11 @@ extern "C" DPCOMP_GPU_RUNTIME_EXPORT void _mlir_ciface_kernel_barrier(int64_t) {
   STUB();
 }
 
+extern "C" DPCOMP_GPU_RUNTIME_EXPORT void
+_mlir_ciface_kernel_mem_fence(int64_t) {
+  STUB();
+}
+
 #define ATOMIC_FUNC_DECL(op, suff, dt)                                         \
   extern "C" DPCOMP_GPU_RUNTIME_EXPORT dt _mlir_ciface_atomic_##op##_##suff(   \
       void *, dt) {                                                            \

mlir/include/mlir-extensions/dialect/gpu_runtime/IR/GpuRuntimeOps.td

@@ -164,5 +164,13 @@ def GPUBarrierOp : GpuRuntime_Op<"barrier"> {
   let assemblyFormat = "$flags attr-dict";
 }
 
+def GPUMemFenceOp : GpuRuntime_Op<"mem_fence"> {
+  let summary = "Orders loads and stores of a work-item executing a kernel.";
+
+  let arguments = (ins GpuRuntime_FenceFlagsAttr:$flags);
+
+  let assemblyFormat = "$flags attr-dict";
+}
+
 #endif // GPURUNTIME_OPS
 

mlir/lib/Conversion/gpu_to_gpu_runtime.cpp

@@ -927,6 +927,32 @@ class ConvertBarrierOp
   }
 };
 
+class ConvertMemFenceOp
+    : public mlir::OpConversionPattern<gpu_runtime::GPUMemFenceOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(gpu_runtime::GPUMemFenceOp op,
+                  gpu_runtime::GPUMemFenceOp::Adaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    auto scope = mlir::spirv::Scope::Workgroup;
+    mlir::spirv::MemorySemantics semantics;
+    if (adaptor.flags() == gpu_runtime::FenceFlags::global) {
+      semantics = mlir::spirv::MemorySemantics::SequentiallyConsistent |
+                  mlir::spirv::MemorySemantics::CrossWorkgroupMemory;
+    } else if (adaptor.flags() == gpu_runtime::FenceFlags::local) {
+      semantics = mlir::spirv::MemorySemantics::SequentiallyConsistent |
+                  mlir::spirv::MemorySemantics::WorkgroupMemory;
+    } else {
+      return mlir::failure();
+    }
+    rewriter.replaceOpWithNewOp<mlir::spirv::MemoryBarrierOp>(op, scope,
+                                                              semantics);
+    return mlir::success();
+  }
+};
+
 // TODO: something better
 class ConvertFunc : public mlir::OpConversionPattern<mlir::FuncOp> {
 public:
@@ -1006,11 +1032,11 @@ struct GPUToSpirvPass
     mlir::arith::populateArithmeticToSPIRVPatterns(typeConverter, patterns);
     mlir::populateMathToSPIRVPatterns(typeConverter, patterns);
 
-    patterns.insert<ConvertSubviewOp, ConvertCastOp<mlir::memref::CastOp>,
-                    ConvertCastOp<mlir::memref::ReinterpretCastOp>,
-                    ConvertLoadOp, ConvertStoreOp, ConvertAtomicOps,
-                    ConvertFunc, ConvertAssert, ConvertBarrierOp>(typeConverter,
-                                                                  context);
+    patterns
+        .insert<ConvertSubviewOp, ConvertCastOp<mlir::memref::CastOp>,
+                ConvertCastOp<mlir::memref::ReinterpretCastOp>, ConvertLoadOp,
+                ConvertStoreOp, ConvertAtomicOps, ConvertFunc, ConvertAssert,
+                ConvertBarrierOp, ConvertMemFenceOp>(typeConverter, context);
 
     if (failed(
             applyFullConversion(kernelModules, *target, std::move(patterns))))

numba_dpcomp/numba_dpcomp/mlir/kernel_impl.py

@@ -320,3 +320,21 @@ def _barrier_impl(builder, flags=None):
 @infer_global(barrier)
 class _BarrierId(ConcreteTemplate):
     cases = [signature(types.void, types.int64), signature(types.void)]
+
+
+def mem_fence(flags=None):
+    _stub_error()
+
+
+@registry.register_func("mem_fence", mem_fence)
+def _memf_fence_impl(builder, flags=None):
+    if flags is None:
+        flags = CLK_GLOBAL_MEM_FENCE
+
+    res = 0  # TODO: remove
+    return builder.external_call("kernel_mem_fence", inputs=flags, outputs=res)
+
+
+@infer_global(mem_fence)
+class _MemFenceId(ConcreteTemplate):
+    cases = [signature(types.void, types.int64), signature(types.void)]

numba_dpcomp/numba_dpcomp/mlir/kernel_sim.py

@@ -34,6 +34,7 @@
     atomic_add,
     atomic_sub,
     barrier,
+    mem_fence,
 )
 
 _ExecutionState = namedtuple(
@@ -82,13 +83,11 @@ def sub(arr, ind, val):
 
 
 def barrier_proxy(flags):
-    global _greenlet_found
-    assert _greenlet_found, "greenlet package not installed"
     state = get_exec_state()
-    assert len(state.tasks) > 0
     wg_size = state.wg_size[0]
     assert wg_size > 0
     if wg_size > 1:
+        assert len(state.tasks) > 0
         indices = copy.deepcopy(state.indices)
         next_task = state.current_task[0] + 1
         if next_task >= wg_size:
@@ -98,6 +97,10 @@ def barrier_proxy(flags):
         state.indices[:] = indices
 
 
+def mem_fence_proxy(flags):
+    pass  # Nothing
+
+
 def _setup_execution_state(global_size, local_size):
     import numba_dpcomp.mlir.kernel_impl
 
@@ -129,6 +132,7 @@ def _destroy_execution_state():
     ("atomic_add", atomic_add, atomic_proxy.add),
     ("atomic_sub", atomic_sub, atomic_proxy.sub),
     ("barrier", barrier, barrier_proxy),
+    ("mem_fence", mem_fence, mem_fence_proxy),
 ]
 
 
@@ -179,6 +183,11 @@ def wrapper():
 _barrier_ops = ["barrier"]
 
 
+def _have_barrier_ops(func):
+    g = func.__globals__
+    return any(n in g for n in _barrier_ops)
+
+
 def _execute_kernel(global_size, local_size, func, *args):
     if len(local_size) == 0:
         local_size = (1,) * len(global_size)
@@ -188,7 +197,7 @@ def _execute_kernel(global_size, local_size, func, *args):
     state = _setup_execution_state(global_size, local_size)
     try:
         groups = tuple((g + l - 1) // l for g, l in zip(global_size, local_size))
-        need_barrier = any(n in func.__globals__ for n in _barrier_ops)
+        need_barrier = max(local_size) > 1 and _have_barrier_ops(func)
         for gid in product(*(range(g) for g in groups)):
             offset = tuple(g * l for g, l in zip(gid, local_size))
             size = tuple(
@@ -202,7 +211,7 @@ def _execute_kernel(global_size, local_size, func, *args):
 
             if need_barrier:
                 global _greenlet_found
-                assert _greenlet_found
+                assert _greenlet_found, "greenlet package not installed"
                 tasks = state.tasks
                 assert len(tasks) == 0
                 for indices in product(*indices_range):

numba_dpcomp/numba_dpcomp/mlir/tests/test_gpu.py

@@ -29,6 +29,7 @@
     kernel_func,
     DEFAULT_LOCAL_SIZE,
     barrier,
+    mem_fence,
     CLK_LOCAL_MEM_FENCE,
     CLK_GLOBAL_MEM_FENCE,
 )
@@ -687,6 +688,38 @@ def func(c):
     assert_equal(gpu_res, sim_res)
 
 
+@require_gpu
+@pytest.mark.parametrize("op", [barrier, mem_fence])
+@pytest.mark.parametrize("flags", [CLK_LOCAL_MEM_FENCE, CLK_GLOBAL_MEM_FENCE])
+@pytest.mark.parametrize("global_size", [1, 2, 27])
+@pytest.mark.parametrize("local_size", [1, 2, 7])
+def test_barrier_ops(op, flags, global_size, local_size):
+    atomic_add = atomic.add
+
+    def func(a, b):
+        i = get_global_id(0)
+        v = a[i]
+        op(flags)
+        b[i] = a[i]
+
+    sim_func = kernel_sim(func)
+    gpu_func = kernel_cached(func)
+
+    a = np.arange(global_size, dtype=np.int64)
+
+    sim_res = np.zeros(global_size, a.dtype)
+    sim_func[global_size, local_size](a, sim_res)
+
+    gpu_res = np.zeros(global_size, a.dtype)
+
+    with print_pass_ir([], ["ConvertParallelLoopToGpu"]):
+        gpu_func[global_size, local_size](a, gpu_res)
+        ir = get_print_buffer()
+        assert ir.count("gpu.launch blocks") == 1, ir
+
+    assert_equal(gpu_res, sim_res)
+
+
 @require_gpu
 @pytest.mark.parametrize("global_size", [1, 2, 27])
 @pytest.mark.parametrize("local_size", [1, 2, 7])

numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/lower_to_gpu.cpp

@@ -1186,7 +1186,9 @@ class ConvertBarrierOps : public mlir::OpRewritePattern<mlir::func::CallOp> {
     using funcptr_t = void (*)(mlir::Operation *, mlir::PatternRewriter &,
                                gpu_runtime::FenceFlags);
     const std::pair<llvm::StringRef, funcptr_t> handlers[] = {
-        {"kernel_barrier", &genBarrierOp<gpu_runtime::GPUBarrierOp>}};
+        {"kernel_barrier", &genBarrierOp<gpu_runtime::GPUBarrierOp>},
+        {"kernel_mem_fence", &genBarrierOp<gpu_runtime::GPUMemFenceOp>},
+    };
 
     auto funcName = op.getCallee();
     for (auto &h : handlers) {