[python][gpy] Kernel api: local mem allocation support (#202)

Author: Hardcode84

dpcomp_gpu_runtime/lib/kernel_api_stubs.cpp

@@ -73,3 +73,24 @@ ATOMIC_FUNC_DECL2(sub)
 
 #undef ATOMIC_FUNC_DECL2
 #undef ATOMIC_FUNC_DECL
+
+#define LOCAL_ARRAY_FUNC_DECL(type, cnt)                                       \
+  extern "C" DPCOMP_GPU_RUNTIME_EXPORT void                                    \
+      _mlir_ciface_local_array_##type##_##cnt() {                              \
+    STUB();                                                                    \
+  }
+
+#define LOCAL_ARRAY_FUNC_DECL2(cnt)                                            \
+  LOCAL_ARRAY_FUNC_DECL(int32, cnt)                                            \
+  LOCAL_ARRAY_FUNC_DECL(int64, cnt)                                            \
+  LOCAL_ARRAY_FUNC_DECL(float32, cnt)                                          \
+  LOCAL_ARRAY_FUNC_DECL(float64, cnt)
+
+LOCAL_ARRAY_FUNC_DECL2(1)
+LOCAL_ARRAY_FUNC_DECL2(2)
+LOCAL_ARRAY_FUNC_DECL2(3)
+LOCAL_ARRAY_FUNC_DECL2(4)
+LOCAL_ARRAY_FUNC_DECL2(5)
+
+#undef LOCAL_ARRAY_FUNC_DECL2
+#undef LOCAL_ARRAY_FUNC_DECL

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

@@ -50,7 +50,17 @@ def GpuRuntime_FenceFlags : I32EnumAttr<"FenceFlags",
   let genSpecializedAttr = 0;
   let cppNamespace = "::gpu_runtime";
 }
-def GpuRuntime_FenceFlagsAttr : EnumAttr<GpuRuntime_Dialect, GpuRuntime_FenceFlags, "fenceFlags">;
+def GpuRuntime_FenceFlagsAttr : EnumAttr<GpuRuntime_Dialect, GpuRuntime_FenceFlags, "fence_flags">;
+
+def GpuRuntime_StorageClass : I32EnumAttr<"StorageClass",
+    "Kernel barrier and fence flags",
+    [
+      I32EnumAttrCase<"local", 1>,
+    ]>{
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::gpu_runtime";
+}
+def GpuRuntime_StorageClassAttr : EnumAttr<GpuRuntime_Dialect, GpuRuntime_StorageClass, "storage_class">;
 
 def CreateGpuStreamOp : GpuRuntime_Op<"create_gpu_stream", [NoSideEffect]> {
   let results = (outs GpuRuntime_OpaqueType : $result);

mlir/lib/Conversion/gpu_to_gpu_runtime.cpp

@@ -212,10 +212,11 @@ struct InsertGPUAllocs
                     if (op->getDialect() == scfDialect ||
                         mlir::isa<mlir::ViewLikeOpInterface>(op))
                       continue;
-                    if (mlir::dyn_cast<mlir::memref::AllocOp>(op)) {
-                      gpuBufferAllocs.insert({op, {}});
-                    } else if (mlir::dyn_cast<mlir::memref::GetGlobalOp>(op)) {
+                    if (mlir::isa<mlir::memref::AllocOp,
+                                  mlir::memref::GetGlobalOp>(op)) {
                       gpuBufferAllocs.insert({op, {}});
+                    } else if (mlir::isa<mlir::func::CallOp>(op)) {
+                      // Ignore
                     } else {
                       op->emitError("Unhandled memref producer");
                       return mlir::WalkResult::interrupt();
@@ -942,6 +943,104 @@ class ConvertMemFenceOp
   }
 };
 
+static llvm::Optional<mlir::spirv::StorageClass>
+convertStorageClass(mlir::Attribute src) {
+  auto attr = src.dyn_cast_or_null<gpu_runtime::StorageClassAttr>();
+  if (!attr)
+    return llvm::None;
+
+  auto sc = attr.getValue();
+  if (sc == gpu_runtime::StorageClass::local)
+    return mlir::spirv::StorageClass::Workgroup;
+
+  return llvm::None;
+}
+
+static mlir::spirv::StorageClass
+convertStorageClass(mlir::Attribute src, mlir::spirv::StorageClass def) {
+  auto ret = convertStorageClass(src);
+  if (ret)
+    return *ret;
+
+  return def;
+}
+
+class ConvertGlobalOp
+    : public mlir::OpConversionPattern<mlir::memref::GlobalOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::memref::GlobalOp op,
+                  mlir::memref::GlobalOp::Adaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    auto memrefType = op.type();
+    if (!memrefType.hasStaticShape())
+      return mlir::failure();
+
+    auto storageClass = convertStorageClass(memrefType.getMemorySpace());
+    if (!storageClass)
+      return mlir::failure();
+
+    auto converter = getTypeConverter();
+    assert(converter);
+
+    auto elemType = converter->convertType(memrefType.getElementType());
+    if (!elemType)
+      return mlir::failure();
+
+    auto elemCount = memrefType.getNumElements();
+    auto newType = mlir::spirv::ArrayType::get(elemType, elemCount);
+    auto ptrType = mlir::spirv::PointerType::get(newType, *storageClass);
+
+    rewriter.replaceOpWithNewOp<mlir::spirv::GlobalVariableOp>(
+        op, ptrType, adaptor.sym_name());
+    return mlir::success();
+  }
+};
+
+class ConvertGetGlobalOp
+    : public mlir::OpConversionPattern<mlir::memref::GetGlobalOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::memref::GetGlobalOp op,
+                  mlir::memref::GetGlobalOp::Adaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    auto memrefType = op.getType().dyn_cast<mlir::MemRefType>();
+    if (!memrefType)
+      return mlir::failure();
+
+    auto storageClass = convertStorageClass(memrefType.getMemorySpace());
+    if (!storageClass)
+      return mlir::failure();
+
+    auto converter = getTypeConverter();
+    assert(converter);
+    auto resType = converter->convertType(memrefType);
+    if (!resType)
+      return mlir::failure();
+
+    auto elemType = converter->convertType(memrefType.getElementType());
+    if (!elemType)
+      return mlir::failure();
+
+    auto elemCount = memrefType.getNumElements();
+    auto newType = mlir::spirv::ArrayType::get(elemType, elemCount);
+    auto ptrType = mlir::spirv::PointerType::get(newType, *storageClass);
+
+    auto loc = op->getLoc();
+    mlir::Value res =
+        rewriter.create<mlir::spirv::AddressOfOp>(loc, ptrType, adaptor.name());
+    if (res.getType() != resType)
+      res = rewriter.create<mlir::spirv::BitcastOp>(loc, resType, res);
+
+    rewriter.replaceOp(op, res);
+    return mlir::success();
+  }
+};
+
 // TODO: something better
 class ConvertFunc : public mlir::OpConversionPattern<mlir::FuncOp> {
 public:
@@ -1024,12 +1123,18 @@ struct GPUToSpirvPass
     mlir::RewritePatternSet patterns(context);
 
     typeConverter.addConversion(
-        [](mlir::MemRefType type) -> llvm::Optional<mlir::Type> {
-          if (type.hasRank() && type.getElementType().isIntOrFloat())
-            return mlir::spirv::PointerType::get(
-                type.getElementType(),
-                mlir::spirv::StorageClass::CrossWorkgroup);
-          return mlir::Type(nullptr);
+        [&typeConverter](mlir::MemRefType type) -> llvm::Optional<mlir::Type> {
+          if (!type.hasRank() || !type.getElementType().isIntOrFloat())
+            return mlir::Type(nullptr);
+
+          auto elemType = typeConverter.convertType(type.getElementType());
+          if (!elemType)
+            return mlir::Type(nullptr);
+
+          auto sc = convertStorageClass(
+              type.getMemorySpace(), mlir::spirv::StorageClass::CrossWorkgroup);
+
+          return mlir::spirv::PointerType::get(elemType, sc);
         });
 
     mlir::ScfToSPIRVContext scfToSpirvCtx;
@@ -1044,8 +1149,8 @@ struct GPUToSpirvPass
         .insert<ConvertSubviewOp, ConvertCastOp<mlir::memref::CastOp>,
                 ConvertCastOp<mlir::memref::ReinterpretCastOp>, ConvertLoadOp,
                 ConvertStoreOp, ConvertAtomicOps, ConvertFunc, ConvertAssert,
-                ConvertBarrierOp, ConvertMemFenceOp, ConvertUndef>(
-            typeConverter, context);
+                ConvertBarrierOp, ConvertMemFenceOp, ConvertUndef,
+                ConvertGlobalOp, ConvertGetGlobalOp>(typeConverter, context);
 
     if (failed(
             applyFullConversion(kernelModules, *target, std::move(patterns))))

mlir/lib/dialect/plier_util/dialect.cpp

@@ -1180,7 +1180,7 @@ struct SignCastCastPropagate : public mlir::OpRewritePattern<CastOp> {
     if (!signCast)
       return mlir::failure();
 
-    auto srcType = op.source().getType().template cast<mlir::ShapedType>();
+    auto srcType = signCast.getType().template cast<mlir::ShapedType>();
     auto dstType = op.getType().template cast<mlir::ShapedType>();
     if (srcType.getElementType() != dstType.getElementType() ||
         !srcType.hasRank() || !dstType.hasRank())
@@ -1200,6 +1200,88 @@ struct SignCastCastPropagate : public mlir::OpRewritePattern<CastOp> {
   }
 };
 
+struct SignCastReinterpretPropagate
+    : public mlir::OpRewritePattern<mlir::memref::ReinterpretCastOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::memref::ReinterpretCastOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    auto signCast = op.source().getDefiningOp<plier::SignCastOp>();
+    if (!signCast)
+      return mlir::failure();
+
+    auto srcType = signCast.getType().cast<mlir::ShapedType>();
+    auto dstType = op.getType().cast<mlir::MemRefType>();
+    if (srcType.getElementType() != dstType.getElementType())
+      return mlir::failure();
+
+    auto src = signCast.value();
+    auto finalType = src.getType().cast<mlir::MemRefType>();
+
+    auto newDstType =
+        mlir::MemRefType::get(dstType.getShape(), dstType.getElementType(),
+                              dstType.getLayout(), finalType.getMemorySpace());
+
+    auto loc = op.getLoc();
+    auto offset = op.getMixedOffsets().front();
+    auto sizes = op.getMixedSizes();
+    auto strides = op.getMixedStrides();
+    auto cast = rewriter.createOrFold<mlir::memref::ReinterpretCastOp>(
+        loc, newDstType, src, offset, sizes, strides);
+    rewriter.replaceOpWithNewOp<plier::SignCastOp>(op, dstType, cast);
+
+    return mlir::success();
+  }
+};
+
+struct SignCastLoadPropagate
+    : public mlir::OpRewritePattern<mlir::memref::LoadOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::memref::LoadOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    auto signCast = op.memref().getDefiningOp<plier::SignCastOp>();
+    if (!signCast)
+      return mlir::failure();
+
+    auto loc = op.getLoc();
+    auto src = signCast.value();
+    auto newOp =
+        rewriter.createOrFold<mlir::memref::LoadOp>(loc, src, op.indices());
+
+    if (newOp.getType() != op.getType())
+      newOp = rewriter.create<plier::SignCastOp>(loc, op.getType(), newOp);
+
+    rewriter.replaceOp(op, newOp);
+    return mlir::success();
+  }
+};
+
+struct SignCastStorePropagate
+    : public mlir::OpRewritePattern<mlir::memref::StoreOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::memref::StoreOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    auto signCast = op.memref().getDefiningOp<plier::SignCastOp>();
+    if (!signCast)
+      return mlir::failure();
+
+    auto src = signCast.value();
+    auto srcElemType = src.getType().cast<mlir::MemRefType>().getElementType();
+    auto val = op.value();
+    if (val.getType() != srcElemType)
+      val = rewriter.create<plier::SignCastOp>(op.getLoc(), srcElemType, val);
+
+    rewriter.replaceOpWithNewOp<mlir::memref::StoreOp>(op, val, src,
+                                                       op.indices());
+    return mlir::success();
+  }
+};
+
 template <typename Op>
 struct SignCastAllocPropagate
     : public mlir::OpRewritePattern<plier::SignCastOp> {
@@ -1223,7 +1305,7 @@ struct SignCastAllocPropagate
 
 struct SignCastTensorFromElementsPropagate
     : public mlir::OpRewritePattern<plier::SignCastOp> {
-  using mlir::OpRewritePattern<plier::SignCastOp>::OpRewritePattern;
+  using OpRewritePattern::OpRewritePattern;
 
   mlir::LogicalResult
   matchAndRewrite(plier::SignCastOp op,
@@ -1422,7 +1504,8 @@ void SignCastOp::getCanonicalizationPatterns(::mlir::RewritePatternSet &results,
       SignCastCastPropagate<mlir::tensor::CastOp>,
       SignCastCastPropagate<mlir::memref::CastOp>,
       SignCastCastPropagate<plier::ChangeLayoutOp>,
-      SignCastAllocPropagate<mlir::memref::AllocOp>,
+      SignCastReinterpretPropagate, SignCastLoadPropagate,
+      SignCastStorePropagate, SignCastAllocPropagate<mlir::memref::AllocOp>,
       SignCastAllocPropagate<mlir::memref::AllocaOp>,
       SignCastTensorFromElementsPropagate, SignCastTensorCollapseShapePropagate,
       SignCastBuferizationPropagate<mlir::bufferization::ToMemrefOp>,

numba_dpcomp/numba_dpcomp/mlir/kernel_impl.py

@@ -16,6 +16,8 @@
 
 from numba import prange
 from numba.core import types
+from numba.core.typing.npydecl import parse_dtype, parse_shape
+from numba.core.types.npytypes import Array
 from numba.core.typing.templates import (
     AbstractTemplate,
     ConcreteTemplate,
@@ -338,3 +340,40 @@ def _memf_fence_impl(builder, flags=None):
 @infer_global(mem_fence)
 class _MemFenceId(ConcreteTemplate):
     cases = [signature(types.void, types.int64), signature(types.void)]
+
+
+class local(Stub):
+    pass
+
+
+def local_array(shape, dtype):
+    _stub_error()
+
+
+setattr(local, "array", local_array)
+
+
+@infer_global(local_array)
+class _LocalId(AbstractTemplate):
+    def generic(self, args, kws):
+        shape = kws["shape"] if "shape" in kws else args[0]
+        dtype = kws["dtype"] if "dtype" in kws else args[1]
+
+        ndim = parse_shape(shape)
+        dtype = parse_dtype(dtype)
+        arr_type = Array(dtype=dtype, ndim=ndim, layout="C")
+        return signature(arr_type, shape, dtype)
+
+
+@registry.register_func("local_array", local_array)
+def _local_array_impl(builder, shape, dtype):
+    try:
+        len(shape)  # will raise if not available
+    except:
+        shape = (shape,)
+
+    func_name = f"local_array_{dtype_str(builder, dtype)}_{len(shape)}"
+    res = builder.init_tensor(shape, dtype)
+    return builder.external_call(
+        func_name, inputs=shape, outputs=res, return_tensor=True
+    )