Generate spirv::BuiltIn::GlobalInvocationId (#170)

Author: Hardcode84

mlir/include/mlir-extensions/dialect/plier_util/CMakeLists.txt

@@ -4,6 +4,8 @@ set(dialect_namespace plier_util)
 set(LLVM_TARGET_DEFINITIONS ${dialect}.td)
 mlir_tablegen(${dialect}Enums.h.inc -gen-enum-decls)
 mlir_tablegen(${dialect}Enums.cpp.inc -gen-enum-defs)
+mlir_tablegen(${dialect}Attributes.h.inc -gen-attrdef-decls -attrdefs-dialect=plier_util)
+mlir_tablegen(${dialect}Attributes.cpp.inc -gen-attrdef-defs -attrdefs-dialect=plier_util)
 mlir_tablegen(${dialect}.h.inc -gen-op-decls)
 mlir_tablegen(${dialect}.cpp.inc -gen-op-defs)
 mlir_tablegen(${dialect}Dialect.h.inc -gen-dialect-decls -dialect=${dialect_namespace})

mlir/include/mlir-extensions/dialect/plier_util/PlierUtilOps.td

@@ -17,6 +17,7 @@
 
 include "mlir/IR/OpBase.td"
 include "mlir/Dialect/GPU/GPUBase.td"
+include "mlir/IR/EnumAttr.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
@@ -274,4 +275,26 @@ def ReleaseContextOp : PlierUtil_Op<"release_context"> {
   let arguments = (ins PlierUtil_OpaqueType:$context);
 }
 
+def GPUEx_Dimension : I32EnumAttr<"GpuDimension",
+    "a dimension, either 'x', 'y', or 'z'",
+    [
+      I32EnumAttrCase<"x", 0>,
+      I32EnumAttrCase<"y", 1>,
+      I32EnumAttrCase<"z", 2>
+    ]>{
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::plier";
+}
+def GPUEx_DimensionAttr : EnumAttr<PlierUtil_Dialect, GPUEx_Dimension, "dim">;
+
+class GPU_IndexOp<string mnemonic, list<Trait> traits = []> :
+    PlierUtil_Op<mnemonic, !listconcat(traits, [NoSideEffect])>,
+    Arguments<(ins GPUEx_DimensionAttr:$dimension)>, Results<(outs Index)> {
+  let assemblyFormat = "$dimension attr-dict";
+}
+
+def GPU_GlobalIdOp : GPU_IndexOp<"gpu_global_id"> {
+  let hasCanonicalizer = 1;
+}
+
 #endif // PLIER_UTIL_OPS

mlir/include/mlir-extensions/dialect/plier_util/dialect.hpp

@@ -27,7 +27,11 @@
 #include <mlir/Dialect/GPU/GPUDialect.h>
 
 #include "mlir-extensions/dialect/plier_util/PlierUtilOpsDialect.h.inc"
-//#include "mlir-extensions/dialect/plier_util/PlierUtilOpsEnums.h.inc"
+#include "mlir-extensions/dialect/plier_util/PlierUtilOpsEnums.h.inc"
+
+#define GET_ATTRDEF_CLASSES
+#include "mlir-extensions/dialect/plier_util/PlierUtilOpsAttributes.h.inc"
+
 #define GET_OP_CLASSES
 #include "mlir-extensions/dialect/plier_util/PlierUtilOps.h.inc"
 

mlir/lib/dialect/plier_util/dialect.cpp

@@ -56,8 +56,14 @@ void PlierUtilDialect::initialize() {
 #define GET_OP_LIST
 #include "mlir-extensions/dialect/plier_util/PlierUtilOps.cpp.inc"
       >();
+
   addTypes<OpaqueType>();
   addInterfaces<PlierUtilInlinerInterface>();
+
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "mlir-extensions/dialect/plier_util/PlierUtilOpsAttributes.cpp.inc"
+      >();
 }
 
 mlir::Type PlierUtilDialect::parseType(mlir::DialectAsmParser &parser) const {
@@ -1508,11 +1514,64 @@ void TakeContextOp::build(mlir::OpBuilder &b, mlir::OperationState &result,
   build(b, result, allTypes, initFunc, releaseFunc);
 }
 
+namespace {
+struct GenGlobalId : public mlir::OpRewritePattern<mlir::arith::AddIOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::arith::AddIOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+
+    auto getArg = [](auto op, bool rev) -> mlir::Value {
+      return rev ? op.getLhs() : op.getRhs();
+    };
+
+    mlir::gpu::Dimension dim;
+    mlir::arith::MulIOp other;
+    for (auto rev : {false, true}) {
+      auto arg1 = getArg(op, rev);
+      auto arg2 = getArg(op, !rev);
+      if (auto tid = arg1.getDefiningOp<mlir::gpu::ThreadIdOp>()) {
+        dim = tid.dimension();
+        other = arg2.getDefiningOp<mlir::arith::MulIOp>();
+        break;
+      }
+    }
+
+    if (!other)
+      return mlir::failure();
+
+    for (auto rev : {false, true}) {
+      auto arg1 = getArg(other, rev).getDefiningOp<mlir::gpu::BlockIdOp>();
+      auto arg2 = getArg(other, !rev).getDefiningOp<mlir::gpu::BlockDimOp>();
+      if (arg1 && arg2) {
+        if (arg1.dimension() != dim || arg2.dimension() != dim)
+          return mlir::failure();
+
+        rewriter.replaceOpWithNewOp<plier::GlobalIdOp>(
+            op, static_cast<plier::GpuDimension>(dim));
+        return mlir::success();
+      }
+    }
+
+    return mlir::failure();
+  }
+};
+} // namespace
+
+void GlobalIdOp::getCanonicalizationPatterns(::mlir::RewritePatternSet &results,
+                                             ::mlir::MLIRContext *context) {
+  results.insert<GenGlobalId>(context);
+}
+
 } // namespace plier
 
 #include "mlir-extensions/dialect/plier_util/PlierUtilOpsDialect.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir-extensions/dialect/plier_util/PlierUtilOps.cpp.inc"
 
-//#include "mlir-extensions/dialect/plier_util/PlierUtilOpsEnums.cpp.inc"
+#define GET_ATTRDEF_CLASSES
+#include "mlir-extensions/dialect/plier_util/PlierUtilOpsAttributes.cpp.inc"
+
+#include "mlir-extensions/dialect/plier_util/PlierUtilOpsEnums.cpp.inc"

numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/lower_to_gpu.cpp

@@ -870,10 +870,9 @@ struct AbiAttrsPass
   void runOnOperation() override {
     auto gpuModule = getOperation();
     auto *context = &getContext();
-    auto attrName = mlir::spirv::getEntryPointABIAttrName();
-    // TODO: Check if block size is const and set appropriate.
-    const int32_t sizes[] = {0, 0, 0};
-    auto abi = mlir::spirv::getEntryPointABIAttr(sizes, context);
+    auto attrName =
+        mlir::StringAttr::get(context, mlir::spirv::getEntryPointABIAttrName());
+    auto abi = mlir::spirv::getEntryPointABIAttr(llvm::None, context);
     for (auto gpuFunc : gpuModule.getOps<mlir::gpu::GPUFuncOp>()) {
       if (!mlir::gpu::GPUDialect::isKernel(gpuFunc) ||
           gpuFunc->getAttr(attrName))
@@ -1295,7 +1294,7 @@ struct GPUToSpirvPass
                                                                context);
 
     patterns.insert<LaunchConfigConversion<
-        mlir::gpu::BlockDimOp, mlir::spirv::BuiltIn::WorkgroupSize>>(
+        plier::GlobalIdOp, mlir::spirv::BuiltIn::GlobalInvocationId>>(
         typeConverter, context);
 
     if (failed(
@@ -3129,7 +3128,7 @@ static void populateLowerToGPUPipelineLow(mlir::OpPassManager &pm) {
   commonOptPasses(funcPM);
   funcPM.addPass(std::make_unique<KernelMemrefOpsMovementPass>());
   funcPM.addPass(std::make_unique<GpuLaunchSinkOpsPass>());
-  //  funcPM.addPass(std::make_unique<SinkGpuDimsPass>());
+  funcPM.addPass(std::make_unique<SinkGpuDimsPass>());
   pm.addPass(mlir::createGpuKernelOutliningPass());
   pm.addPass(mlir::createSymbolDCEPass());