implement launch with cpp

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
2026-07-01 08:51:56 +08:00 · 2026-01-30 15:05:05 +08:00
parent 37a8f4a5be
commit 524cf83708
10 changed files with 250 additions and 243 deletions
@@ -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