implement launch with cpp

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

include/launch.h

@@ -0,0 +1,55 @@
+#ifndef TRITON_TVM_FFI_LAUNCH_H_
+#define TRITON_TVM_FFI_LAUNCH_H_
+
+#include "type.h"
+#include <tvm/ffi/object.h>
+
+namespace triton_tvm_ffi {
+
+class TVMFFILauncherImplObj : public tvm::ffi::Object {
+public:
+  TVMFFILauncherImplObj(const tvm::ffi::Array<Type> &signature,
+                        bool launchCooperativeGrid, bool launchAsync);
+  TVMFFILauncherImplObj(const TVMFFILauncherImplObj &other) = default;
+  TVMFFILauncherImplObj(TVMFFILauncherImplObj &&other) = default;
+  void Launch(int32_t gridX, int32_t gridY, int32_t gridZ, uint64_t stream,
+              uint64_t function,
+              tvm::ffi::Tuple<int32_t, int32_t, int32_t> kernelMetadata,
+              tvm::ffi::ObjectRef launchMetadata,
+              tvm::ffi::ObjectRef launchEnterHook,
+              tvm::ffi::ObjectRef launchExitHook,
+              tvm::ffi::ObjectRef globalScratchObject,
+              tvm::ffi::ObjectRef profileScratchObject,
+              const tvm::ffi::Array<tvm::ffi::Any> &kernelArgs) const;
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("triton_tvm_ffi.TVMFFILauncherImpl",
+                                    TVMFFILauncherImplObj, tvm::ffi::Object);
+
+private:
+  tvm::ffi::Array<Type> signature_;
+  const bool launchCooperativeGrid_;
+  const bool launchAsync_;
+};
+
+class TVMFFILauncherImpl : public tvm::ffi::ObjectRef {
+public:
+  TVMFFILauncherImpl(tvm::ffi::Array<Type> signature,
+                     bool launchCooperativeGrid, bool launchAsync);
+  using tvm::ffi::ObjectRef::ObjectRef;
+  using tvm::ffi::ObjectRef::operator=;
+  void Launch(int32_t gridX, int32_t gridY, int32_t gridZ, uint64_t stream,
+              uint64_t function,
+              tvm::ffi::Tuple<int32_t, int32_t, int32_t> kernelMetadata,
+              tvm::ffi::ObjectRef launchMetadata,
+              tvm::ffi::ObjectRef launchEnterHook,
+              tvm::ffi::ObjectRef launchExitHook,
+              tvm::ffi::ObjectRef globalScratchObject,
+              tvm::ffi::ObjectRef profileScratchObject,
+              const tvm::ffi::Array<tvm::ffi::Any> &kernelArgs) const;
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TVMFFILauncherImpl,
+                                                tvm::ffi::ObjectRef,
+                                                TVMFFILauncherImplObj);
+};
+
+} // namespace triton_tvm_ffi
+
+#endif

include/macro.h

@@ -1,8 +1,19 @@
 #ifndef TRITON_TVM_FFI_MACRO_H_
 #define TRITON_TVM_FFI_MACRO_H_
 
+#include "exception.h"
+
 #if defined(__GNUC__) || defined(__clang__)
 #define TRITON_TVM_FFI_INLINE __attribute__((always_inline)) inline
 #endif
 
+#define UNLIKELY(cond) __builtin_expect((cond), 0)
+
+#define CUDA_CHECK(code)                                                       \
+  do {                                                                         \
+    if (UNLIKELY((code) != CUDA_SUCCESS)) {                                    \
+      throw triton_tvm_ffi::CUDAException(code);                               \
+    }                                                                          \
+  } while (false)
+
 #endif

include/value.h

@@ -1,52 +0,0 @@
-#ifndef TRITON_TVM_FFI_VALUE_H_
-#define TRITON_TVM_FFI_VALUE_H_
-
-#include "macro.h"
-#include "type.h"
-#include <tvm/ffi/any.h>
-#include <tvm/ffi/object.h>
-
-namespace triton_tvm_ffi {
-
-class TypedValueObj : public tvm::ffi::Object {
-public:
-  TypedValueObj(Type type, const tvm::ffi::Any &value);
-  TypedValueObj(Type type, tvm::ffi::Any &&value);
-  TypedValueObj(const TypedValueObj &other) = default;
-  TypedValueObj(TypedValueObj &&other) = default;
-  TypedValueObj &operator=(const TypedValueObj &other) = default;
-  TypedValueObj &operator=(TypedValueObj &&other) = default;
-  TRITON_TVM_FFI_INLINE Type GetType() const { return type_; }
-  TRITON_TVM_FFI_INLINE const tvm::ffi::Any &GetValue() const { return value_; }
-  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("triton_tvm_ffi.TypedValue", TypedValueObj,
-                                    tvm::ffi::Object);
-
-private:
-  Type type_;
-  tvm::ffi::Any value_;
-};
-
-class TypedValue : public tvm::ffi::ObjectRef {
-public:
-  TypedValue(Type type, const tvm::ffi::Any &value);
-  TypedValue(Type type, tvm::ffi::Any &&value);
-  using tvm::ffi::ObjectRef::ObjectRef;
-  using tvm::ffi::ObjectRef::operator=;
-  TRITON_TVM_FFI_INLINE Type GetType() const { return get()->GetType(); }
-  TRITON_TVM_FFI_INLINE const tvm::ffi::Any &GetValue() const {
-    return get()->GetValue();
-  }
-  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TypedValue, tvm::ffi::ObjectRef,
-                                                TypedValueObj);
-};
-
-tvm::ffi::Optional<TypedValue> MakeTypedValue(const tvm::ffi::String &type,
-                                              const tvm::ffi::Any &value);
-
-tvm::ffi::Array<TypedValue>
-MakeTypedValues(const tvm::ffi::Array<tvm::ffi::String> &types,
-                const tvm::ffi::Array<tvm::ffi::Any> &values);
-
-} // namespace triton_tvm_ffi
-
-#endif