[gpu] Sink UndefOp into gpu kernel (#201)

Author: Hardcode84

mlir/lib/Conversion/gpu_to_gpu_runtime.cpp

@@ -971,6 +971,25 @@ class ConvertAssert : public mlir::OpConversionPattern<mlir::cf::AssertOp> {
   }
 };
 
+class ConvertUndef : public mlir::OpConversionPattern<plier::UndefOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(plier::UndefOp op, plier::UndefOp::Adaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    auto converter = getTypeConverter();
+    assert(converter);
+
+    auto resType = converter->convertType(op.getType());
+    if (!resType)
+      return mlir::failure();
+
+    rewriter.replaceOpWithNewOp<mlir::spirv::UndefOp>(op, resType);
+    return mlir::success();
+  }
+};
+
 struct GPUToSpirvPass
     : public mlir::PassWrapper<GPUToSpirvPass,
                                mlir::OperationPass<mlir::ModuleOp>> {
@@ -1025,7 +1044,8 @@ struct GPUToSpirvPass
         .insert<ConvertSubviewOp, ConvertCastOp<mlir::memref::CastOp>,
                 ConvertCastOp<mlir::memref::ReinterpretCastOp>, ConvertLoadOp,
                 ConvertStoreOp, ConvertAtomicOps, ConvertFunc, ConvertAssert,
-                ConvertBarrierOp, ConvertMemFenceOp>(typeConverter, context);
+                ConvertBarrierOp, ConvertMemFenceOp, ConvertUndef>(
+            typeConverter, context);
 
     if (failed(
             applyFullConversion(kernelModules, *target, std::move(patterns))))

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

@@ -1213,7 +1213,7 @@ class GpuLaunchSinkOpsPass
             auto isSinkingBeneficiary = [](mlir::Operation *op) -> bool {
               return isa<arith::ConstantOp, func::ConstantOp, arith::SelectOp,
                          arith::CmpIOp, arith::IndexCastOp, arith::MulIOp,
-                         arith::SubIOp, arith::AddIOp>(op);
+                         arith::SubIOp, arith::AddIOp, plier::UndefOp>(op);
             };
 
             // Pull in instructions that can be sunk