[Mlir-commits] [clang] [llvm] [mlir] [AMDGPU] Implement AsyncMark Stages (PR #220442)

[Krzysztof Drewniak]

================
@@ -0,0 +1,205 @@
+//===- AMDGPUAsyncStages.h --------------------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file
+/// Shared AMDGPU asyncmark stage definitions.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_SUPPORT_AMDGPUASYNCSTAGES_H
+#define LLVM_SUPPORT_AMDGPUASYNCSTAGES_H
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/bit.h"
+#include "llvm/ADT/iterator.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/MathExtras.h"
+#include "llvm/Support/raw_ostream.h"
+#include <cassert>
+#include <cstdint>
+#include <iterator>
+
+namespace llvm {
+namespace AMDGPU {
+
+/// Bit mask of the async stages tracked by the asyncmark / wait_asyncmark
+/// intrinsics. Each stage has its own independent sequence of marks.
+///
+/// A single stage is represented as a mask with exactly one bit set. Iterating
+/// over a mask visits each stage it names, from the least significant bit up.
+///
+/// This type also provides certain stronger guarantees than a simple integer:
+///   - Default constructor initializes the mask to zero.
+///   - Constructor ensures undefined bits cannot be set.
+class AsyncStages {
+public:
+  using value_type = uint32_t;
+
+  static constexpr unsigned NUM_STAGES = 11;
+
+  // The bit of each stage is part of the IR, so do not renumber them. Some
+  // values are RESERVED for later use.
+  enum : value_type {
+    NONE = 0,
+    // Tensor loads to LDS and tensor stores from LDS.
+    TENSOR = 1u << 0,
----------------
krzysz00 wrote:

Hopefully we won't need to split this

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