chore: import upstream snapshot with attribution

paddle/cinn/backends/codegen_gpu_dev.h

@@ -0,0 +1,172 @@
+// Copyright (c) 2024 CINN Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include "paddle/cinn/backends/codegen_c.h"
+#include "paddle/cinn/common/common.h"
+#include "paddle/cinn/common/ir_util.h"
+#include "paddle/cinn/ir/ir.h"
+#include "paddle/cinn/ir/ir_printer.h"
+#include "paddle/cinn/ir/lowered_func.h"
+#include "paddle/cinn/ir/module.h"
+#include "paddle/cinn/lang/packed_func.h"
+#include "paddle/cinn/runtime/cinn_runtime.h"
+
+namespace cinn::ir {
+class Module;
+}  // namespace cinn::ir
+
+namespace cinn {
+namespace backends {
+
+/**
+ * CUDA/HIP device code generator.
+ *
+ * It generates the device function, e.g, the function called "myadd" will have
+ * a __global__ function called "myadd_kernel", different from codegen_c, the
+ * declaration of the "myadd_kernel" function has an expanded argument list,
+ * which finally similar to `__global__ void myadd(float* __restrict__ A, float*
+ * __restrict__ B, int n);`
+ */
+class CodeGenGpuDev : public CodeGenC {
+ public:
+  explicit CodeGenGpuDev(Target target);
+
+  /**
+   * Compile the \p module to \p outputs.
+   */
+  void Compile(const ir::Module& module, const Outputs& outputs);
+
+  //! Compile on RTC (RunTime Compilation).
+  std::string Compile(const ir::Module& module, bool use_rtc = true);
+
+  void Compile(const ir::LoweredFunc& func);
+
+  /**
+   * \brief Print a function argument in cuda/hip syntax. Currently, just some
+   * decoration of __restrict__.
+   * @param arg the argument.
+   * @return the representation in cuda/hip syntax.
+   *
+   * We make it a static to make the test easier.
+   */
+  virtual void PrintFuncArg(const ir::Argument& arg);
+
+  std::string Compile(const ir::Module& module, OutputKind output_kind);
+
+  ir::Expr GetDynSharedMemOffset() const {
+    if (MathEqual(dyn_shared_mem_offset_, Expr(-1))) {
+      return Expr(0);
+    }
+    return dyn_shared_mem_offset_;
+  }
+
+ protected:
+  void Visit(const ir::_Var_* op) override;
+  void Visit(const ir::_LoweredFunc_* op) override;
+  void Visit(const ir::Min* op) override;
+  void Visit(const ir::Max* op) override;
+  void Visit(const ir::Call* op) override;
+  void Visit(const ir::Load* op) override;
+  void VisitStmt(const ir::stmt::Free& stmt) override;
+  void VisitStmt(const ir::stmt::Alloc& stmt) override;
+  void VisitStmt(const ir::stmt::Store& op) override;
+  void VisitStmt(const ir::stmt::Let& stmt) override;
+
+  // Print element access at a cuda/hip built-in vector on a load/store node
+  bool PrintBuiltinVectorAccess(const ir::LoadStoreAddrMnger* op,
+                                ir::Expr index,
+                                bool is_store);
+  // Print element access at a cuda/hip built-in vector on a store node
+  bool PrintBuiltinVectorAccess(const ir::stmt::Store& stmt,
+                                ir::Expr index,
+                                bool is_store);
+
+  void PrintBuiltinCodes();
+
+  virtual void PrintIncludes() = 0;
+
+  virtual void PrintTempBufferCreation(const ir::Buffer& buffer);
+
+  void PrintTempBufferAliasDefinition(const ir::Buffer& buffer);
+
+  std::vector<ir::stmt::StmtRef> GenerateBufferAliasStmts(
+      const ir::_LoweredFunc_* op, const std::vector<ir::Buffer>& temp_buffers);
+
+  std::vector<ir::stmt::StmtRef> FilterDeallocTempBuffers(
+      const std::vector<ir::stmt::StmtRef>& frees);
+
+  /**
+   * Print the function declaration, this is different from C, we expand the
+   * arguments and get something like
+   * `__global__ void myadd(float* __restrict__ A, float* __restrict__ B, int
+   * n);`
+   */
+  virtual void PrintFunctionDeclaration(const ir::_LoweredFunc_* op);
+
+  inline void ProcessMinMaxOperand(ir::Expr* a,
+                                   ir::Expr* b,
+                                   int unify_bit,
+                                   bool both_dyn) {
+    if (unify_bit > 0) {
+      std::string type_func = "int" + std::to_string(unify_bit) + "_t";
+      if (both_dyn) {
+        // if both contains dynamic symbol, like: min(S0, S1), it it likely that
+        // S0 is int and S1 is int64_t. So we need to enforce the type cast by
+        // ir::Call
+        *a = ir::Call::Make(common::Int(unify_bit),
+                            type_func,
+                            {*a},
+                            {},
+                            ir::CallType::Intrinsic);
+        *b = ir::Call::Make(common::Int(unify_bit),
+                            type_func,
+                            {*b},
+                            {},
+                            ir::CallType::Intrinsic);
+      } else {
+        *a = ir::Cast::Make(common::Int(unify_bit), *a);
+        *b = ir::Cast::Make(common::Int(unify_bit), *b);
+      }
+    }
+  }
+
+  std::unordered_map<std::string, common::Type>& DynamicShapeMap() {
+    return dynamic_shape_map_;
+  }
+
+ private:
+  Target target_;
+  bool use_rtc_{false};
+  // names of vectorized tensors from `Let` statements where dtypes of the
+  // tensors are customized_type with customized_type::k_builtin_vector_t
+  // prefix
+  std::unordered_set<std::string> vectorized_tensor_names_;
+
+  // This map is used to store the name and type of the dynamic shape func args
+  // so that during codegen of ir::Min & ir::Max call, we can have correct type
+  std::unordered_map<std::string, common::Type> dynamic_shape_map_;
+
+  ir::Expr dyn_shared_mem_offset_{-1};
+  std::vector<ir::Buffer> dynamic_alloc_buffers_;
+};
+
+ir::Expr CalculateSharedMemory(const ir::LoweredFunc& func);
+
+}  // namespace backends
+}  // namespace cinn