Add target description query and verifier pass

Author: zhczhong

include/gc/Analysis/TargetDescriptionAnalysis.h

@@ -0,0 +1,92 @@
+//===-- TargetDescriptionAnalysis.h - target description class --*- C++ -*-===//
+//
+// This file is licensed 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_TARGETDESCRIPTIONANALYSIS_H
+#define MLIR_ANALYSIS_TARGETDESCRIPTIONANALYSIS_H
+
+#include "gc/Dialect/Linalgx/LinalgxOps.h"
+#include "mlir/Dialect/DLTI/DLTI.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "llvm/ADT/StringRef.h"
+
+namespace mlir {
+namespace gc {
+
+using namespace mlir;
+
+enum DeviceType { CPU = 0 };
+
+class TargetDescriptionAnalysisBase {
+public:
+  TargetDescriptionAnalysisBase(Operation *op, DeviceType device)
+      : ctx(op->getContext()), device(device),
+        layout(isa<ModuleOp>(op) ? dyn_cast<ModuleOp>(op)
+                                 : op->getParentOfType<ModuleOp>()),
+        loc(op->getLoc()) {}
+
+  // get the device ID
+  DeviceType getDevice() { return device; }
+
+  // get the MLIR context
+  MLIRContext *getContext() { return ctx; }
+
+  // get the data layout
+  DataLayout getLayout() { return layout; }
+
+  // get the property value by key
+  std::optional<Attribute> getPropertyValue(StringRef key);
+
+  // get the location
+  Location getLocation() { return loc; }
+
+  // check if the property exists
+  bool hasProperty(StringRef key) { return getPropertyValue(key).has_value(); }
+
+  // emit warning if the property is not found
+  template <typename T>
+  void emitNotFoundWarning(Location loc, StringRef key, T value);
+
+  // the map from device type to device string
+  static llvm::DenseMap<DeviceType, std::string> DeviceKeyMap;
+
+private:
+  MLIRContext *ctx;
+  DeviceType device;
+  DataLayout layout;
+  Location loc;
+};
+
+class CPUTargetDescriptionAnalysis : public TargetDescriptionAnalysisBase {
+public:
+  static constexpr StringLiteral kL1CacheSize = "L1_cache_size_in_bytes";
+  static constexpr StringLiteral kL2CacheSize = "L2_cache_size_in_bytes";
+  static constexpr StringLiteral kL3CacheSize = "L3_cache_size_in_bytes";
+  static constexpr StringLiteral kMaxVectorWidth = "max_vector_width";
+  static constexpr StringLiteral kNumThreads = "num_threads";
+
+  // get runtime OMP_NUM_THREADS
+  size_t getNumThreads();
+
+  // get cache size by cacheLevel
+  size_t getCacheSize(uint8_t cacheLevel);
+
+  // get the maximum vector length in bits
+  size_t getMaxVectorWidth();
+
+  // get the default value map(attr key, default value)
+  static llvm::DenseMap<StringRef, int64_t> CPUTargetDeafultValueMap;
+
+  CPUTargetDescriptionAnalysis(Operation *op)
+      : TargetDescriptionAnalysisBase(op, DeviceType::CPU) {}
+};
+
+} // namespace gc
+} // namespace mlir
+
+#endif

include/gc/Transforms/Passes.td

@@ -19,46 +19,47 @@ def TileLinalgNamed : Pass<"tile-named-linalg", "func::FuncOp"> {
 
 #ifdef GC_HAS_ONEDNN_DIALECT
 def ConvertOneDNNGraphToLinalg : Pass<"convert-onednn-graph-to-linalg"> {
-  let summary = "Lower the operations from the oneDNN Graph dialect into Linalg";
-  let description = [{
-    Lowers the `onednn_graph` ops to `linalg` ops.
-  }];
+  let summary =
+      "Lower the operations from the oneDNN Graph dialect into Linalg";
+  let description = [{Lowers the `onednn_graph` ops to `linalg` ops.}];
   let dependentDialects = [
-      "func::FuncDialect",
-      "math::MathDialect",
-      "arith::ArithDialect",
-      "tensor::TensorDialect",
-      "linalg::LinalgDialect",
-      "linalgx::LinalgxDialect"
+    "func::FuncDialect", "math::MathDialect", "arith::ArithDialect",
+    "tensor::TensorDialect", "linalg::LinalgDialect", "linalgx::LinalgxDialect"
   ];
 }
 #endif
 
 #ifdef GC_USE_IMEX
 def LinalgToXeGPU : Pass<"linalg-to-xegpu", "func::FuncOp"> {
   let summary = "Convert linalg dialect to XeGPU dialect.";
-  let description = [{
-    Lower linalg ops to XeGPU dialect.
-  }];
-  let dependentDialects = ["linalg::LinalgDialect",
-                           "gpu::GPUDialect",
-                           "xegpu::XeGPUDialect",
-                           "scf::SCFDialect",
-                           "memref::MemRefDialect",
-                           "arith::ArithDialect",
-                           "math::MathDialect",
-                           "vector::VectorDialect"];
+  let description = [{Lower linalg ops to XeGPU dialect.}];
+  let dependentDialects = [
+    "linalg::LinalgDialect", "gpu::GPUDialect", "xegpu::XeGPUDialect",
+    "scf::SCFDialect", "memref::MemRefDialect", "arith::ArithDialect",
+    "math::MathDialect", "vector::VectorDialect"
+  ];
   let options = [
     Option<"kTile", "k-tile", "int64_t",
-           /*default=*/"32",
-           "GEMM tile size for reduction dimension.">,
+           /*default=*/"32", "GEMM tile size for reduction dimension.">,
     Option<"stages", "stages", "int64_t",
-           /*default=*/"1",
-           "Number of cooperative prefetch stages.">,
+           /*default=*/"1", "Number of cooperative prefetch stages.">,
     ListOption<"dpasTile", "dpas-tile", "int64_t",
                "DPAS register block sizes MxNxK">,
   ];
 }
 #endif
 
+def VerifyTargetDescription : Pass<"verify-target-description", "ModuleOp"> {
+  let summary = "Verify the target description from ModuleOp DLTI attribute.";
+  let description = [{
+    Verify the target description from ModuleOp DLTI attribute. Raise error for unexpected input(such as a negative number of num_threads), and raise warn for missing fields, and provide a default value(such as 32K for L1_cache_size).
+  }];
+  let dependentDialects = ["DLTIDialect"];
+  let options = [
+    Option<"device", "device", "std::string",
+           /*default=*/"\"CPU\"",
+           "The device to verify. Supported device: CPU, ">,
+  ];
+}
+
 #endif // GC_DIALECT_GC_PASSES

lib/gc/Analysis/CMakeLists.txt

@@ -0,0 +1,15 @@
+gc_set_mlir_link_components(MLIR_LINK_COMPONENTS
+  MLIRIR
+  MLIRSupport)
+
+gc_add_mlir_library(GCAnalysis
+  TargetDescriptionAnalysis.cpp
+
+  DEPENDS
+    GraphCompilerPassIncGen
+
+  LINK_LIBS PUBLIC
+    ${mlir_dialect_libs}
+    ${MLIR_LINK_COMPONENTS}
+    GcInterface
+  )

lib/gc/Analysis/TargetDescriptionAnalysis.cpp

@@ -0,0 +1,98 @@
+//===-- TargetDescriptionAnalysis.cpp - target description impl -*- C++ -*-===//
+//
+// This file is licensed 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 "gc/Analysis/TargetDescriptionAnalysis.h"
+#include <limits>
+#include <llvm/Support/Debug.h>
+
+namespace mlir {
+namespace gc {
+
+#define DEBUG_TYPE "target-description-analysis"
+
+llvm::DenseMap<DeviceType, std::string>
+    TargetDescriptionAnalysisBase::DeviceKeyMap = {
+        {CPU, "CPU"},
+};
+
+// default values for properties
+llvm::DenseMap<StringRef, int64_t>
+    CPUTargetDescriptionAnalysis::CPUTargetDeafultValueMap = {
+        {CPUTargetDescriptionAnalysis::kNumThreads, 1},
+        {CPUTargetDescriptionAnalysis::kL1CacheSize, 32 * 1024},
+        {CPUTargetDescriptionAnalysis::kL2CacheSize, 1024 * 1024},
+        {CPUTargetDescriptionAnalysis::kL3CacheSize, 32 * 1024 * 1024},
+        {CPUTargetDescriptionAnalysis::kMaxVectorWidth, 512},
+};
+
+template <typename T>
+void TargetDescriptionAnalysisBase::emitNotFoundWarning(Location loc,
+                                                        StringRef key,
+                                                        T value) {
+  mlir::emitWarning(loc) << key << " not found, using default value " << value;
+}
+
+static int64_t getIntFromAttribute(Attribute attr) {
+  if (auto intAttr = dyn_cast<IntegerAttr>(attr)) {
+    if (intAttr.getType().isSignedInteger())
+      return intAttr.getSInt();
+    else if (intAttr.getType().isUnsignedInteger())
+      return intAttr.getUInt();
+    else
+      return intAttr.getInt();
+  }
+  llvm_unreachable("Not an integer attribute");
+}
+
+std::optional<Attribute>
+TargetDescriptionAnalysisBase::getPropertyValue(StringRef key) {
+  return layout.getDevicePropertyValue(
+      Builder(getContext())
+          .getStringAttr(DeviceKeyMap[getDevice()] /* device ID*/),
+      Builder(getContext()).getStringAttr(key));
+}
+
+size_t CPUTargetDescriptionAnalysis::getNumThreads() {
+  std::optional<Attribute> numThreads = getPropertyValue(kNumThreads);
+
+  if (numThreads && isa<IntegerAttr>(*numThreads))
+    return getIntFromAttribute(*numThreads);
+  emitNotFoundWarning(getLocation(), kNumThreads,
+                      CPUTargetDeafultValueMap[kNumThreads]);
+  return CPUTargetDeafultValueMap[kNumThreads];
+}
+
+size_t CPUTargetDescriptionAnalysis::getCacheSize(uint8_t cacheLevel) {
+  assert(cacheLevel > 0 && cacheLevel < 4 && "Invalid cache level");
+  StringLiteral key = "";
+  if (cacheLevel == 1)
+    key = kL1CacheSize;
+  else if (cacheLevel == 2)
+    key = kL2CacheSize;
+  else if (cacheLevel == 3)
+    key = kL3CacheSize;
+
+  std::optional<Attribute> cacheSize = getPropertyValue(key);
+  if (cacheSize && isa<IntegerAttr>(*cacheSize))
+    return getIntFromAttribute(*cacheSize);
+
+  emitNotFoundWarning(getLocation(), key, CPUTargetDeafultValueMap[key]);
+  return CPUTargetDeafultValueMap[key];
+}
+
+size_t CPUTargetDescriptionAnalysis::getMaxVectorWidth() {
+  std::optional<Attribute> maxVectorWidth = getPropertyValue(kMaxVectorWidth);
+  if (maxVectorWidth && isa<IntegerAttr>(*maxVectorWidth))
+    return getIntFromAttribute(*maxVectorWidth);
+  emitNotFoundWarning(getLocation(), kMaxVectorWidth,
+                      CPUTargetDeafultValueMap[kMaxVectorWidth]);
+  return CPUTargetDeafultValueMap[kMaxVectorWidth];
+}
+
+} // namespace gc
+} // namespace mlir

lib/gc/CAPI/CMakeLists.txt

@@ -1,6 +1,7 @@
 set(GC_ALL_LIBS
   ${GC_ONEDNN_DIALECT_LIB_NAME}
   GcPasses
+  GCAnalysis
   MLIRCPURuntimeTransforms)
 
 if(GC_ENABLE_IMEX)

lib/gc/CMakeLists.txt

@@ -1,3 +1,4 @@
+add_subdirectory(Analysis)
 add_subdirectory(CAPI)
 add_subdirectory(Dialect)
 add_subdirectory(Transforms)

lib/gc/ExecutionEngine/Driver/CMakeLists.txt

@@ -39,4 +39,5 @@ gc_add_mlir_library(GcJitWrapper
     ${dialect_libs}
     ${conversion_libs}
     ${GC_PASSES}
+    GCAnalysis
   )

lib/gc/Transforms/CMakeLists.txt

@@ -13,6 +13,7 @@ gc_add_mlir_library(GcPasses
   OneDNNGraphToLinalg.cpp
   Pipeline.cpp
   TileNamed.cpp
+  VerifyTargetDescription.cpp
 
   DEPENDS
     GraphCompilerPassIncGen