[mlir][spirv] Convert `ub.poison` to `spirv.undef`

[Hardcode84]

Summary

SPIR-V doesn't have poison, but poison can be converted to undef.

[kuhar]

Thanks for taking care of this, @Hardcode84!

mlir/include/mlir/Conversion/Passes.td

1064

Nit: 'Converts supported UB ops to SPIR-V dialect ops.'

mlir/lib/Conversion/UBToSPIRV/UBToSPIRV.cpp

25

39–41

Could you define this inline? There no benefit to outlining these matchAndRewirte functions. This may be inconsistent with the older SPIR-V passed, but we converged on in-line definition in newer ones.

42

Shouldn't we use the type from the adaptor, or are these the same? Also, if the adaptor is unused, the argument should be unnamed (, OpAdaptor, )

62–63

mlir/test/Conversion/UBToSPIRV/ub-to-spirv.mlir

15–16

Could you add one test with vector<4xi32> to document how the non-scalar case is handled?

[Hardcode84]

mlir/lib/Conversion/UBToSPIRV/UBToSPIRV.cpp

42

I don't see return type on adaptor, and I want to access unconverted type anyway, to avoid converting user types, which can coincidentally also be converted to primitive types.

[kuhar]

mlir/lib/Conversion/UBToSPIRV/UBToSPIRV.cpp

42

SG, thanks for the explanation

[Hardcode84]

mlir/test/Conversion/UBToSPIRV/ub-to-spirv.mlir

15–16

vector is not covered yet (will be rejected by isIntOrIndexOrFloat)

[Hardcode84]

review comments

[kuhar]

mlir/test/Conversion/UBToSPIRV/ub-to-spirv.mlir

15–16

This doesn't prevent us from having a test with a TODO for unsupported types that checks that at least the code doesn't crash

[kuhar]

LGTM % testcase for an unsupported type

[Hardcode84]

vector negative test

mlir/include/mlir/Conversion/Passes.h

 #include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
 #include "mlir/Conversion/TensorToLinalg/TensorToLinalgPass.h"
 #include "mlir/Conversion/TensorToSPIRV/TensorToSPIRVPass.h"
 #include "mlir/Conversion/TosaToArith/TosaToArith.h"
 #include "mlir/Conversion/TosaToLinalg/TosaToLinalg.h"
 #include "mlir/Conversion/TosaToSCF/TosaToSCF.h"
 #include "mlir/Conversion/TosaToTensor/TosaToTensor.h"
 #include "mlir/Conversion/UBToLLVM/UBToLLVM.h"
+#include "mlir/Conversion/UBToSPIRV/UBToSPIRV.h"
 #include "mlir/Conversion/VectorToArmSME/VectorToArmSME.h"
 #include "mlir/Conversion/VectorToGPU/VectorToGPU.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Conversion/VectorToSPIRV/VectorToSPIRVPass.h"
 
 namespace mlir {
 
 /// Generate the code for registering conversion passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Conversion/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_PASSES_H

mlir/include/mlir/Conversion/Passes.td

def UBToLLVMConversionPass : Pass<"convert-ub-to-llvm"> {
   let options = [
     Option<"indexBitwidth", "index-bitwidth", "unsigned",
            /*default=kDeriveIndexBitwidthFromDataLayout*/"0",
            "Bitwidth of the index type, 0 to use size of machine word">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
+// UBToSPIRV
+//===----------------------------------------------------------------------===//
+
+def UBToSPIRVConversionPass : Pass<"convert-ub-to-spirv"> {
+  let summary = "Convert UB dialect to SPIR-V dialect";
+  let description = [{
+    This pass converts supported UB ops to SPIR-V dialect instructions.
+  }];
+  let dependentDialects = ["spirv::SPIRVDialect"];
+}
+
+//===----------------------------------------------------------------------===//
 // VectorToGPU
 //===----------------------------------------------------------------------===//
 
 def ConvertVectorToGPU : Pass<"convert-vector-to-gpu"> {
   let summary = "Lower the operations from the vector dialect into the GPU "
                 "dialect";
   let constructor = "mlir::createConvertVectorToGPUPass()";
   let dependentDialects = [

mlir/include/mlir/Conversion/UBToSPIRV/UBToSPIRV.h

+//===- UBToSPIRV.h - UB to SPIR-V dialect conversion ------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_UBTOSPIRV_UBSPIRV_H
+#define MLIR_CONVERSION_UBTOSPIRV_UBSPIRV_H
+
+#include <memory>
+
+namespace mlir {
+
+class SPIRVTypeConverter;
+class RewritePatternSet;
+class Pass;
+
+#define GEN_PASS_DECL_UBTOSPIRVCONVERSIONPASS
+#include "mlir/Conversion/Passes.h.inc"
+
+namespace ub {
+void populateUBToSPIRVConversionPatterns(SPIRVTypeConverter &converter,
+                                         RewritePatternSet &patterns);
+} // namespace ub
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_UBTOSPIRV_UBSPIRV_H

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(SPIRVToLLVM)
 add_subdirectory(TensorToLinalg)
 add_subdirectory(TensorToSPIRV)
 add_subdirectory(TosaToArith)
 add_subdirectory(TosaToLinalg)
 add_subdirectory(TosaToSCF)
 add_subdirectory(TosaToTensor)
 add_subdirectory(UBToLLVM)
+add_subdirectory(UBToSPIRV)
 add_subdirectory(VectorToArmSME)
 add_subdirectory(VectorToGPU)
 add_subdirectory(VectorToLLVM)
 add_subdirectory(VectorToSCF)
 add_subdirectory(VectorToSPIRV)

mlir/lib/Conversion/UBToSPIRV/CMakeLists.txt

+add_mlir_conversion_library(MLIRUBToSPIRV
+  UBToSPIRV.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/UBToSPIRV
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRSPIRVConversion
+  MLIRSPIRVDialect
+  MLIRUBDialect
+  )

mlir/lib/Conversion/UBToSPIRV/UBToSPIRV.cpp

+//===- UBToSPIRV.cpp - UB to SPIRV-V dialect conversion -------------------===//
+//
+// 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 "mlir/Conversion/UBToSPIRV/UBToSPIRV.h"
+
+#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_UBTOSPIRVCONVERSIONPASS
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+
+struct PoisonOpLowering : public OpConversionPattern<ub::PoisonOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ub::PoisonOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// PoisonOpLowering
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+PoisonOpLowering::matchAndRewrite(ub::PoisonOp op, OpAdaptor adaptor,
+                                  ConversionPatternRewriter &rewriter) const {
+  Type origType = op.getType();
+  if (!origType.isIntOrIndexOrFloat())
+    return rewriter.notifyMatchFailure(
+        op, [&](Diagnostic &diag) { diag << "unsupported type " << origType; });
+
+  Type resType = getTypeConverter()->convertType(origType);
+  if (!resType)
+    return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) {
+      diag << "failed to convert result type " << origType;
+    });
+
+  rewriter.replaceOpWithNewOp<spirv::UndefOp>(op, resType);
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// Pass Definition
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct UBToSPIRVConversionPass
+    : public impl::UBToSPIRVConversionPassBase<UBToSPIRVConversionPass> {
+  using Base::Base;
+
+  void runOnOperation() override {
+    Operation *op = getOperation();
+    spirv::TargetEnvAttr targetAttr = spirv::lookupTargetEnvOrDefault(op);
+    std::unique_ptr<SPIRVConversionTarget> target =
+        SPIRVConversionTarget::get(targetAttr);
+
+    SPIRVConversionOptions options;
+    SPIRVTypeConverter typeConverter(targetAttr, options);
+
+    RewritePatternSet patterns(&getContext());
+    ub::populateUBToSPIRVConversionPatterns(typeConverter, patterns);
+
+    if (failed(applyPartialConversion(op, *target, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Pattern Population
+//===----------------------------------------------------------------------===//
+
+void mlir::ub::populateUBToSPIRVConversionPatterns(
+    SPIRVTypeConverter &converter, RewritePatternSet &patterns) {
+  patterns.add<PoisonOpLowering>(converter, patterns.getContext());
+}

mlir/test/Conversion/UBToSPIRV/ub-to-spirv.mlir

+// RUN: mlir-opt -split-input-file -convert-ub-to-spirv -verify-diagnostics %s | FileCheck %s
+
+module attributes {
+  spirv.target_env = #spirv.target_env<
+    #spirv.vce<v1.0, [Int8, Int16, Int64, Float16, Float64, Shader], []>, #spirv.resource_limits<>>
+} {
+
+// CHECK-LABEL: @check_poison
+func.func @check_poison(%lhs: i32, %rhs: i32) {
+// CHECK: {{.*}} = spirv.Undef : i32
+  %0 = ub.poison : index
+// CHECK: {{.*}} = spirv.Undef : i16
+  %1 = ub.poison : i16
+// CHECK: {{.*}} = spirv.Undef : f64
+  %2 = ub.poison : f64
+  return
+}
+
+}