Try fix ck decoder compilation in fbcode (2/n) (#1012)

* fix includes in decoder and splitk decoder * resolve compilation error by adding -Werror to cpp extensions compiler flags and fixing a hidden override
facebookresearch · Apr 2, 2024 · 09773f8 · 09773f8
1 parent 7fffd3d
commit 09773f8
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 8 deletions.
diff --git a/setup.py b/setup.py
@@ -377,6 +377,8 @@ def get_extensions():
                 "-U__CUDA_NO_HALF_CONVERSIONS__",
                 "-DCK_FMHA_FWD_FAST_EXP2=1",
                 "-fgpu-flush-denormals-to-zero",
+                "-Werror",
+                "-Woverloaded-virtual",
             ]
             + generator_flag
             + cc_flag,

diff --git a/xformers/csrc/attention/hip_fmha/attention_forward_decoder.cpp b/xformers/csrc/attention/hip_fmha/attention_forward_decoder.cpp
@@ -149,7 +149,7 @@ at::Tensor& efficient_attention_forward_decoder_ck_out_impl(
             lds_bytes);
 
         auto invoker = device_op_t::Invoker{};
-        (void)invoker.Run(arg, {stream});
+        (void)invoker.Run(&arg, {stream});
       });
 
   return O;

diff --git a/xformers/csrc/attention/hip_fmha/attention_forward_splitk.cpp b/xformers/csrc/attention/hip_fmha/attention_forward_splitk.cpp
@@ -159,7 +159,7 @@ at::Tensor& efficient_attention_forward_decoder_splitk_ck_out_impl(
             lds_bytes);
 
         auto invoker = device_op_t::Invoker{};
-        (void)invoker.Run(arg, {stream});
+        (void)invoker.Run(&arg, {stream});
       });
 
   return O;

diff --git a/xformers/csrc/attention/hip_fmha/ck_attention_forward_decoder.h b/xformers/csrc/attention/hip_fmha/ck_attention_forward_decoder.h
@@ -432,9 +432,11 @@ struct FMHADecoderSeqlen1DeviceOp : public BaseOperator {
 
   struct Invoker : public BaseInvoker {
     using Argument = DeviceOp::Argument;
-    float Run(
-        const Argument& arg,
-        const StreamConfig& stream_config = StreamConfig{}) {
+    virtual float Run(
+        const BaseArgument* base_arg,
+        const StreamConfig& stream_config = StreamConfig{}) override {
+      // copy so it's alive while being used
+      auto arg = *dynamic_cast<const Argument*>(base_arg);
       auto threads_per_wavefront = arg.block_dim.x;
 
       auto Q_size_k_alignment_necessary = 0;

diff --git a/xformers/csrc/attention/hip_fmha/ck_attention_forward_decoder_splitk.h b/xformers/csrc/attention/hip_fmha/ck_attention_forward_decoder_splitk.h
@@ -586,9 +586,12 @@ struct FMHADecoderSplitKDeviceOp : public BaseOperator {
 
   struct Invoker : public BaseInvoker {
     using Argument = DeviceOp::Argument;
-    float Run(
-        const Argument& arg,
-        const StreamConfig& stream_config = StreamConfig{}) {
+    virtual float Run(
+        const BaseArgument* base_arg,
+        const StreamConfig& stream_config = StreamConfig{}) override {
+      // copy so it's alive while being used
+      auto arg = *dynamic_cast<const Argument*>(base_arg);
+
       auto threads_per_wavefront = arg.block_dim.x;
       auto Q_size_k_alignment_necessary = 0;