Implemented dedicated executor for functions with static shapes

Author: AndreyPavlenko

include/gc/ExecutionEngine/GPURuntime/GpuOclRuntime.h

@@ -128,7 +128,8 @@ struct OclRuntime {
   struct Exports;
   friend OclContext;
   friend OclModuleBuilder;
-  template <unsigned N> friend struct OclModuleExecutor;
+  template <unsigned N> friend struct DynamicExecutor;
+  template <unsigned N> friend struct StaticExecutor;
   explicit OclRuntime(const Ext &ext);
   const Ext &ext;
 
@@ -173,12 +174,13 @@ struct OclContext {
   OclContext(const OclContext &) = delete;
   OclContext &operator=(const OclContext &) = delete;
 
-  void finish();
+  [[nodiscard]] llvm::Expected<bool> finish();
 
 private:
   friend OclRuntime;
   friend OclRuntime::Exports;
-  template <unsigned N> friend struct OclModuleExecutor;
+  template <unsigned N> friend struct DynamicExecutor;
+  template <unsigned N> friend struct StaticExecutor;
   std::unordered_set<void *> *clPtrs;
 
   void setLastEvent(cl_event event) {
@@ -195,6 +197,9 @@ struct OclContext {
 
 struct OclModule {
   const OclRuntime runtime;
+  // If all the function arguments have static shapes, then this field is true
+  // and main.staticMain is used. Otherwise, main.wrappedMain is used.
+  const bool isStatic;
 
   ~OclModule();
   OclModule(const OclModule &) = delete;
@@ -204,24 +209,35 @@ struct OclModule {
 
 private:
   friend OclModuleBuilder;
-  template <unsigned N> friend struct OclModuleExecutor;
-  using MainFunc = void (*)(void **);
+  template <unsigned N> friend struct DynamicExecutor;
+  template <unsigned N> friend struct OclModuleExecutorBase;
+  template <unsigned N> friend struct StaticExecutor;
+  // This function is only created when all args are memrefs with static shape.
+  using StaticMainFunc = void (*)(OclContext *, void **);
+  // Wrapper, generated by the engine. The arguments are pointers to the values.
+  using WrappedMainFunc = void (*)(void **);
+  union MainFunc {
+    StaticMainFunc staticMain;
+    WrappedMainFunc wrappedMain;
+  };
   const MainFunc main;
   const FunctionType functionType;
   std::unique_ptr<ExecutionEngine> engine;
 
-  explicit OclModule(const OclRuntime &runtime, const MainFunc main,
-                     func::FuncOp functionOp,
+  explicit OclModule(const OclRuntime &runtime, const bool isStatic,
+                     const MainFunc main, const FunctionType functionType,
                      std::unique_ptr<ExecutionEngine> engine)
-      : runtime(runtime), main(main),
-        functionType(functionOp.getFunctionType()), engine(std::move(engine)) {}
+      : runtime(runtime), isStatic(isStatic), main(main),
+        functionType(functionType), engine(std::move(engine)) {}
 };
 
 struct OclModuleBuilder {
   friend OclRuntime;
   explicit OclModuleBuilder(ModuleOp module);
   explicit OclModuleBuilder(OwningOpRef<ModuleOp> &module)
       : OclModuleBuilder(module.release()) {}
+  explicit OclModuleBuilder(OwningOpRef<ModuleOp> &&module)
+      : OclModuleBuilder(module.release()) {}
 
   llvm::Expected<std::shared_ptr<const OclModule>>
   build(const OclRuntime &runtime);
@@ -243,105 +259,141 @@ struct OclModuleBuilder {
   build(const OclRuntime::Ext &ext);
 };
 
-// The main function arguments are added in the following format -
-// https://mlir.llvm.org/docs/TargetLLVMIR/#c-compatible-wrapper-emission.
 // NOTE: This class is mutable and not thread-safe!
-// NOTE: The argument values are not copied, only the pointers are stored!
-template <unsigned N = 64> struct OclModuleExecutor {
-  explicit OclModuleExecutor(std::shared_ptr<const OclModule> &mod)
+template <unsigned N> struct OclModuleExecutorBase {
+
+  void reset() {
+    args.clear();
+    clPtrs.clear();
+    argCounter = 0;
+  }
+
+  Type getArgType(unsigned idx) const {
+    assert(idx < mod->functionType.getNumInputs());
+    return mod->functionType.getInput(idx);
+  }
+
+  [[nodiscard]] bool isSmall() const { return args.small(); }
+
+protected:
+  struct Args : SmallVector<void *, N> {
+    [[nodiscard]] bool small() const { return this->isSmall(); }
+  };
+
+  const std::shared_ptr<const OclModule> &mod;
+  // Contains the pointers of all non-USM arguments. It's expected, that the
+  // arguments are either USM or CL pointers and most probably are USM, thus,
+  // in most cases, this set will be empty.
+  std::unordered_set<void *> clPtrs;
+  Args args;
+  unsigned argCounter = 0;
+
+  explicit OclModuleExecutorBase(std::shared_ptr<const OclModule> &mod)
       : mod(mod) {}
-  OclModuleExecutor(const OclModuleExecutor &) = delete;
-  OclModuleExecutor &operator=(const OclModuleExecutor &) = delete;
-  OclModuleExecutor(const OclModuleExecutor &&) = delete;
-  OclModuleExecutor &operator=(const OclModuleExecutor &&) = delete;
 
-  void exec(OclContext &ctx) {
 #ifndef NDEBUG
+  void checkCtx(const OclContext &ctx) const {
     auto rt = OclRuntime::get(ctx.queue);
     assert(rt);
     assert(*rt == mod->runtime);
+    assert(argCounter == mod->functionType.getNumInputs());
+  }
+
+  void checkArg(void *alignedPtr, bool isUsm = true) const {
+    assert(!isUsm || mod->runtime.isUsm(alignedPtr));
+    // It's recommended to have at least 16-byte alignment
+    assert(reinterpret_cast<std::uintptr_t>(alignedPtr) % 16 == 0);
+  }
 #endif
-    auto size = args.size();
-    auto ctxPtr = &ctx;
-    ctx.clPtrs = &clPtrs;
-    args.emplace_back(&ctxPtr);
-    args.emplace_back(&ctxPtr);
-    args.emplace_back(ZERO_PTR);
-    mod->main(args.data());
-    args.truncate(size);
+};
+
+// NOTE: This executor can only be used if mod->isStatic == true!
+template <unsigned N = 8> struct StaticExecutor : OclModuleExecutorBase<N> {
+  explicit StaticExecutor(std::shared_ptr<const OclModule> &mod)
+      : OclModuleExecutorBase<N>(mod) {
+    assert(this->mod->isStatic);
   }
 
-  void operator()(OclContext &ctx) { exec(ctx); }
+  void exec(OclContext &ctx) {
+#ifndef NDEBUG
+    this->checkCtx(ctx);
+#endif
+    ctx.clPtrs = &this->clPtrs;
+    this->mod->main.staticMain(&ctx, this->args.data());
+  }
 
-  template <typename T>
-  [[nodiscard]] bool operator()(OclContext &ctx, T **ptr1, ...) {
-    {
-      SmallVector<int64_t> values;
-      auto argTypes = mod->functionType.getInputs();
-      unsigned numValues = 0;
-
-      for (unsigned i = 0, n = argTypes.size() - 1; i < n; i++) {
-        if (auto type = llvm::dyn_cast<MemRefType>(argTypes[i])) {
-          if (type.hasStaticShape()) {
-            numValues += type.getShape().size() * 2 + 1;
-            continue;
-          }
-        }
+  void arg(void *alignedPtr, bool isUsm = true) {
 #ifndef NDEBUG
-        OclRuntime::debug(
-            __FILE__, __LINE__,
-            "Only memref arguments with static shape are supported.");
+    this->checkArg(alignedPtr, isUsm);
+    std::ostringstream oss;
+    oss << "Arg" << this->argCounter << ": alignedPtr=" << alignedPtr
+        << ", isUsm=" << (isUsm ? "true" : "false");
+    OclRuntime::debug(__FILE__, __LINE__, oss.str().c_str());
 #endif
-        return false;
-      }
+    ++this->argCounter;
+    this->args.emplace_back(alignedPtr);
+    if (!isUsm) {
+      this->clPtrs.insert(alignedPtr);
+    }
+  }
+
+  template <typename T> void arg(T *alignedPtr, bool isUsm = true) {
+    arg(reinterpret_cast<void *>(alignedPtr), isUsm);
+  }
+
+  void operator()(OclContext &ctx) { exec(ctx); }
 
-      values.reserve(numValues);
-      SmallVector<int64_t> strides;
-      int64_t offset;
+  template <typename T> void operator()(OclContext &ctx, T *ptr1, ...) {
+    {
+      this->reset();
+      arg(reinterpret_cast<void *>(ptr1));
       va_list args;
       va_start(args, ptr1);
-
-      for (unsigned i = 0, n = argTypes.size() - 1; i < n; i++) {
-        auto type = llvm::dyn_cast<MemRefType>(argTypes[i]);
-        strides.clear();
-        if (failed(getStridesAndOffset(type, strides, offset))) {
-#ifndef NDEBUG
-          OclRuntime::debug(__FILE__, __LINE__,
-                            "Failed to get strides and offset.");
-#endif
-          return false;
-        }
-        auto offsetPtr = values.end();
-        values.emplace_back(offset);
-        auto shapePtr = values.end();
-        auto shape = type.getShape();
-        values.append(shape.begin(), shape.end());
-        auto stridesPtr = values.end();
-        values.append(strides.begin(), strides.end());
-        auto ptr =
-            (i == 0) ? reinterpret_cast<void **>(ptr1) : va_arg(args, void **);
-        addArg(*ptr, *ptr, *offsetPtr, shape.size(), shapePtr, stridesPtr);
+      for (unsigned i = 0, n = this->mod->functionType.getNumInputs() - 1;
+           i < n; i++) {
+        arg(va_arg(args, void *));
       }
-
       va_end(args);
       exec(ctx);
-      return true;
     }
   }
+};
 
-  void addArg(void *&alignedPtr, size_t rank, const int64_t *shape,
-              const int64_t *strides, bool isUsm = true) {
-    addArg(alignedPtr, alignedPtr, ZERO, rank, shape, strides, isUsm);
+// The main function arguments are added in the following format -
+// https://mlir.llvm.org/docs/TargetLLVMIR/#c-compatible-wrapper-emission.
+// NOTE: This executor can only be used if mod->isStatic != true!
+template <unsigned N = 64> struct DynamicExecutor : OclModuleExecutorBase<N> {
+  explicit DynamicExecutor(std::shared_ptr<const OclModule> &mod)
+      : OclModuleExecutorBase<N>(mod) {
+    assert(!this->mod->isStatic);
   }
 
-  void addArg(void *&allocatedPtr, void *&alignedPtr, const int64_t &offset,
-              size_t rank, const int64_t *shape, const int64_t *strides,
-              bool isUsm = true) {
+  void exec(OclContext &ctx) {
 #ifndef NDEBUG
-    assert(!isUsm || mod->runtime.isUsm(alignedPtr));
-    // It's recommended to have at least 16-byte alignment
-    assert(reinterpret_cast<std::uintptr_t>(alignedPtr) % 16 == 0);
-    if (auto type = llvm::dyn_cast<MemRefType>(getArgType(argCounter))) {
+    this->checkCtx(ctx);
+#endif
+    auto size = this->args.size();
+    auto ctxPtr = &ctx;
+    this->args.emplace_back(&ctxPtr);
+    this->args.emplace_back(&ctxPtr);
+    this->args.emplace_back(ZERO_PTR);
+    this->mod->main.wrappedMain(this->args.data());
+    this->args.truncate(size);
+  }
+
+  void arg(void *&alignedPtr, size_t rank, const int64_t *shape,
+           const int64_t *strides, bool isUsm = true) {
+    arg(alignedPtr, alignedPtr, ZERO, rank, shape, strides, isUsm);
+  }
+
+  // NOTE: The argument values are not copied, only the pointers are stored!
+  void arg(void *&allocatedPtr, void *&alignedPtr, const int64_t &offset,
+           size_t rank, const int64_t *shape, const int64_t *strides,
+           bool isUsm = true) {
+#ifndef NDEBUG
+    this->checkArg(alignedPtr, isUsm);
+    if (auto type =
+            llvm::dyn_cast<MemRefType>(this->getArgType(this->argCounter))) {
       if (type.hasStaticShape()) {
         auto size = type.getShape();
         assert(rank == size.size());
@@ -361,8 +413,9 @@ template <unsigned N = 64> struct OclModuleExecutor {
     }
 
     std::ostringstream oss;
-    oss << "Arg" << argCounter << ": ptr=" << allocatedPtr
-        << ", alignedPtr=" << alignedPtr << ", offset=" << offset
+    oss << "Arg" << this->argCounter << ": ptr=" << allocatedPtr
+        << ", alignedPtr=" << alignedPtr
+        << ", isUsm=" << (isUsm ? "true" : "false") << ", offset=" << offset
         << ", shape=[";
     for (unsigned i = 0; i < rank; i++) {
       oss << shape[i] << (i + 1 < rank ? ", " : "]");
@@ -374,55 +427,36 @@ template <unsigned N = 64> struct OclModuleExecutor {
     OclRuntime::debug(__FILE__, __LINE__, oss.str().c_str());
 #endif
 
-    argCounter++;
-    args.emplace_back(&allocatedPtr);
-    args.emplace_back(&alignedPtr);
-    args.emplace_back(const_cast<int64_t *>(&offset));
+    ++this->argCounter;
+    this->args.emplace_back(&allocatedPtr);
+    this->args.emplace_back(&alignedPtr);
+    this->args.emplace_back(const_cast<int64_t *>(&offset));
     for (size_t i = 0; i < rank; i++) {
-      args.emplace_back(const_cast<int64_t *>(&shape[i]));
+      this->args.emplace_back(const_cast<int64_t *>(&shape[i]));
     }
     for (size_t i = 0; i < rank; i++) {
-      args.emplace_back(const_cast<int64_t *>(&strides[i]));
+      this->args.emplace_back(const_cast<int64_t *>(&strides[i]));
     }
     if (!isUsm) {
-      clPtrs.insert(alignedPtr);
+      this->clPtrs.insert(alignedPtr);
     }
   }
 
   template <typename T>
-  void addArg(T *&alignedPtr, size_t rank, const int64_t *shape,
-              const int64_t *strides, bool isUsm = true) {
-    addArg(reinterpret_cast<void *&>(alignedPtr), rank, shape, strides, isUsm);
+  void arg(T *&alignedPtr, size_t rank, const int64_t *shape,
+           const int64_t *strides, bool isUsm = true) {
+    arg(reinterpret_cast<void *&>(alignedPtr), rank, shape, strides, isUsm);
   }
 
   template <typename T>
-  void addArg(T *&allocatedPtr, T *&alignedPtr, const int64_t &offset,
-              size_t rank, const int64_t *shape, const int64_t *strides,
-              bool isUsm = true) {
-    addArg(reinterpret_cast<void *&>(allocatedPtr),
-           reinterpret_cast<void *&>(alignedPtr), offset, rank, shape, strides,
-           isUsm);
-  }
-
-  Type getArgType(unsigned idx) const {
-    assert(idx < mod->functionType.getNumInputs() - 1);
-    return mod->functionType.getInput(idx);
-  }
-
-  void reset() {
-    args.clear();
-    clPtrs.clear();
-    argCounter = 0;
+  void arg(T *&allocatedPtr, T *&alignedPtr, const int64_t &offset, size_t rank,
+           const int64_t *shape, const int64_t *strides, bool isUsm = true) {
+    arg(reinterpret_cast<void *&>(allocatedPtr),
+        reinterpret_cast<void *&>(alignedPtr), offset, rank, shape, strides,
+        isUsm);
   }
 
-private:
-  const std::shared_ptr<const OclModule> &mod;
-  // Contains the pointers of all non-USM arguments. It's expected, that the
-  // arguments are either USM or CL pointers and most probably are USM, thus,
-  // in most cases, this set will be empty.
-  std::unordered_set<void *> clPtrs;
-  SmallVector<void *, N + 3> args;
-  unsigned argCounter = 0;
+  void operator()(OclContext &ctx) { exec(ctx); }
 };
 }; // namespace mlir::gc::gpu
 #else

include/gc/Utils/Error.h

@@ -16,7 +16,7 @@
 #include "llvm/Support/Error.h"
 
 namespace mlir::gc::err {
-#ifdef _NDEBUG
+#ifdef NDEBUG
 #define GC_ERR_LOC_DECL
 #define GC_ERR_LOC_ARGS
 #define GC_ERR_LOC