[flang-commits] [flang] [flang][cuda] Add CUFLaunchAttachAttr pass (PR #174465)

[Zhen Wang via flang-commits]

================
@@ -0,0 +1,68 @@
+//===-- CUFLaunchAttachAttr.cpp -------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Dialect/CUF/CUFDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIROpsSupport.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace fir {
+#define GEN_PASS_DEF_CUFLAUNCHATTACHATTR
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+using namespace mlir;
+
+namespace {
+
+class CUFGPUAttachAttrPattern
+    : public OpRewritePattern<mlir::gpu::LaunchFuncOp> {
+  using OpRewritePattern<mlir::gpu::LaunchFuncOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(mlir::gpu::LaunchFuncOp op,
+                                PatternRewriter &rewriter) const override {
+    op->setAttr(cuf::getProcAttrName(),
+                cuf::ProcAttributeAttr::get(op.getContext(),
+                                            cuf::ProcAttribute::Global));
+    return mlir::success();
+  }
+};
+
+struct CUFLaunchAttachAttr
+    : public fir::impl::CUFLaunchAttachAttrBase<CUFLaunchAttachAttr> {
+
+  void runOnOperation() override {
+    auto *context = &this->getContext();
+
+    mlir::RewritePatternSet patterns(context);
+    patterns.add<CUFGPUAttachAttrPattern>(context);
+
+    mlir::ConversionTarget target(*context);
+    target.addIllegalOp<mlir::gpu::LaunchFuncOp>();
+    target.addDynamicallyLegalOp<mlir::gpu::LaunchFuncOp>(
+        [&](mlir::gpu::LaunchFuncOp op) -> bool {
+          if (op.getKernelName().getValue().contains("_cufk_")) {
----------------
wangzpgi wrote:

Can `_cufk_` be used for non cuf kernels, i.e. by the user?

https://github.com/llvm/llvm-project/pull/174465