[FXML-5071] Lower UB to EmitC

mlir/include/mlir/Conversion/Passes.h

@@ -12,7 +12,6 @@
 #include "mlir/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.h"
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/ArithToAMDGPU/ArithToAMDGPU.h"
-#include "mlir/Conversion/ArithToEmitC/ArithToEmitCPass.h"
 #include "mlir/Conversion/ArithToArmSME/ArithToArmSME.h"
 #include "mlir/Conversion/ArithToEmitC/ArithToEmitCPass.h"
 #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
@@ -71,6 +70,7 @@
 #include "mlir/Conversion/TosaToMLProgram/TosaToMLProgram.h"
 #include "mlir/Conversion/TosaToSCF/TosaToSCF.h"
 #include "mlir/Conversion/TosaToTensor/TosaToTensor.h"
+#include "mlir/Conversion/UBToEmitC/UBToEmitC.h"
 #include "mlir/Conversion/UBToLLVM/UBToLLVM.h"
 #include "mlir/Conversion/UBToSPIRV/UBToSPIRV.h"
 #include "mlir/Conversion/VectorToArmSME/VectorToArmSME.h"

mlir/include/mlir/Conversion/Passes.td

@@ -1210,6 +1210,18 @@ def TosaToTensor : Pass<"tosa-to-tensor"> {
   let constructor = "tosa::createTosaToTensor()";
 }
 
+//===----------------------------------------------------------------------===//
+// UBToEmitC
+//===----------------------------------------------------------------------===//
+
+def ConvertUBToEmitC : Pass<"convert-ub-to-emitc"> {
+  let summary = "Convert UB dialect to EmitC dialect";
+  let description = [{
+    This pass converts supported UB ops to EmitC dialect.
+  }];
+  let dependentDialects = ["emitc::EmitCDialect"];
+}
+
 //===----------------------------------------------------------------------===//
 // UBToLLVM
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Conversion/UBToEmitC/UBToEmitC.h

@@ -0,0 +1,25 @@
+//===- UBToEmitC.h - UB to EmitC 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_UBTOEMITC_UBTOEMITC_H
+#define MLIR_CONVERSION_UBTOEMITC_UBTOEMITC_H
+
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+#define GEN_PASS_DECL_CONVERTUBTOEMITC
+#include "mlir/Conversion/Passes.h.inc"
+
+namespace ub {
+void populateUBToEmitCConversionPatterns(TypeConverter &converter,
+                                         RewritePatternSet &patterns);
+} // namespace ub
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_UBTOEMITC_UBTOEMITC_H

mlir/lib/Conversion/CMakeLists.txt

@@ -60,6 +60,7 @@ add_subdirectory(TosaToLinalg)
 add_subdirectory(TosaToMLProgram)
 add_subdirectory(TosaToSCF)
 add_subdirectory(TosaToTensor)
+add_subdirectory(UBToEmitC)
 add_subdirectory(UBToLLVM)
 add_subdirectory(UBToSPIRV)
 add_subdirectory(VectorToArmSME)

mlir/lib/Conversion/UBToEmitC/CMakeLists.txt

@@ -0,0 +1,17 @@
+add_mlir_conversion_library(MLIRUBToEmitC
+  UBToEmitC.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/UBToEmitC
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRLLVMCommonConversion
+  MLIREmitCDialect
+  MLIRUBDialect
+  )

mlir/lib/Conversion/UBToEmitC/UBToEmitC.cpp

@@ -0,0 +1,79 @@
+//===- UBToEmitC.cpp - UB to EmitC 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/UBToEmitC/UBToEmitC.h"
+
+#include "mlir/Dialect/EmitC/IR/EmitC.h"
+#include "mlir/Dialect/EmitC/Transforms/TypeConversions.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/TypeUtilities.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_CONVERTUBTOEMITC
+#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 {
+    const TypeConverter *converter = getTypeConverter();
+    Type convertedType = converter->convertType(op.getType());
+
+    if (!convertedType)
+      return rewriter.notifyMatchFailure(op.getLoc(), "type conversion failed");
+
+    if (!(emitc::isIntegerIndexOrOpaqueType(convertedType) ||
+          emitc::isSupportedFloatType(convertedType))) {
+      return rewriter.notifyMatchFailure(
+          op.getLoc(), "only scalar poison values can be lowered");
+    }
+
+    // Any constant will be fine to lower a poison op
+    rewriter.replaceOpWithNewOp<emitc::VariableOp>(
+        op, convertedType, emitc::OpaqueAttr::get(op->getContext(), ""));
+    return success();
+  }
+};
+} // namespace
+
+void ub::populateUBToEmitCConversionPatterns(TypeConverter &converter,
+                                             RewritePatternSet &patterns) {
+  MLIRContext *ctx = patterns.getContext();
+  patterns.add<PoisonOpLowering>(converter, ctx);
+}
+
+struct ConvertUBToEmitC : public impl::ConvertUBToEmitCBase<ConvertUBToEmitC> {
+  using Base::Base;
+
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    TypeConverter converter;
+    converter.addConversion([](Type t) { return t; });
+    populateEmitCSizeTTypeConversions(converter);
+
+    ConversionTarget target(getContext());
+    target.addLegalDialect<emitc::EmitCDialect>();
+    target.addIllegalDialect<ub::UBDialect>();
+
+    mlir::ub::populateUBToEmitCConversionPatterns(converter, patterns);
+
+    if (failed(applyPartialConversion(getOperation(), target,
+                                      std::move(patterns))))
+      signalPassFailure();
+  }
+};

mlir/test/Conversion/UBToEmitC/convert-ub-to-emitc-unsupported.mlir

@@ -0,0 +1,23 @@
+// RUN: mlir-opt -convert-ub-to-emitc -split-input-file -verify-diagnostics %s
+
+func.func @poison_memref() {
+  // expected-error @+1 {{failed to legalize operation 'ub.poison'}}
+  %0 = ub.poison : memref<i32>
+  return
+}
+
+// -----
+
+func.func @poison_tensor() {
+  // expected-error @+1 {{failed to legalize operation 'ub.poison'}}
+  %1 = ub.poison : tensor<f32>
+  return
+}
+
+// -----
+
+func.func @poison_vector() {
+  // expected-error @+1 {{failed to legalize operation 'ub.poison'}}
+  %1 = "ub.poison"() {value = #ub.poison} : () -> vector<4xi64>
+  return
+}

mlir/test/Conversion/UBToEmitC/convert-ub-to-emitc.mlir

@@ -0,0 +1,12 @@
+// RUN: mlir-opt -convert-ub-to-emitc %s | FileCheck %s
+
+// CHECK-LABEL: func.func @poison
+func.func @poison() {
+  // CHECK: "emitc.variable"(){{.*}}() -> i32
+  %0 = ub.poison : i32
+  // CHECK: "emitc.variable"(){{.*}}() -> f32
+  %1 = ub.poison : f32
+  // CHECK: "emitc.variable"(){{.*}}() -> !emitc.size_t
+  %2 = ub.poison : index
+  return
+}