[mlir][EmitC] Add Arith to EmitC conversions (llvm#84151)

This adds patterns and a pass to convert the Arith dialect to EmitC. For
now, this covers arithemtic binary ops operating on floating point
types.

It is not checked within the patterns whether the types, such as the
Tensor type, are supported in the respective EmitC operations. If
unsupported types should be converted, the conversion will fail anyway
because no legal EmitC operation can be created. This can clearly be
improved in a follow up, also resulting in better error messages.
Functions for such checks should not solely be used in the conversions
and should also be (re)used in the verifier.

Author: marbre

mlir/include/mlir/Conversion/ArithToEmitC/ArithToEmitC.h

@@ -1,22 +1,20 @@
-//===- ArithToEmitC.h - Convert Arith to EmitC ----------------------------===//
+//===- ArithToEmitC.h - Arith to EmitC Patterns -----------------*- 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_ARITHTOEMITC_ARITHTOEMITC_H
 #define MLIR_CONVERSION_ARITHTOEMITC_ARITHTOEMITC_H
 
-#include "mlir/Pass/Pass.h"
-
 namespace mlir {
 class RewritePatternSet;
+class TypeConverter;
 
-#define GEN_PASS_DECL_ARITHTOEMITCCONVERSIONPASS
-#include "mlir/Conversion/Passes.h.inc"
-
-void populateArithToEmitCConversionPatterns(RewritePatternSet &patterns);
+void populateArithToEmitCPatterns(TypeConverter &typeConverter,
+                                  RewritePatternSet &patterns);
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_ARITHTOEMITC_ARITHTOEMITC_H

mlir/include/mlir/Conversion/ArithToEmitC/ArithToEmitCPass.h

@@ -0,0 +1,21 @@
+//===- ArithToEmitCPass.h - Arith to EmitC Pass -----------------*- 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_ARITHTOEMITC_ARITHTOEMITCPASS_H
+#define MLIR_CONVERSION_ARITHTOEMITC_ARITHTOEMITCPASS_H
+
+#include <memory>
+
+namespace mlir {
+class Pass;
+
+#define GEN_PASS_DECL_CONVERTARITHTOEMITC
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_ARITHTOEMITC_ARITHTOEMITCPASS_H

mlir/include/mlir/Conversion/Passes.h

@@ -12,7 +12,7 @@
 #include "mlir/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.h"
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/ArithToAMDGPU/ArithToAMDGPU.h"
-#include "mlir/Conversion/ArithToEmitC/ArithToEmitC.h"
+#include "mlir/Conversion/ArithToEmitC/ArithToEmitCPass.h"
 #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
 #include "mlir/Conversion/ArithToSPIRV/ArithToSPIRV.h"
 #include "mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h"

mlir/include/mlir/Conversion/Passes.td

@@ -125,17 +125,20 @@ def ArithToAMDGPUConversionPass : Pass<"convert-arith-to-amdgpu"> {
   }];
 
   let dependentDialects = ["amdgpu::AMDGPUDialect", "vector::VectorDialect"];
+
+  let options = [
+    Option<"saturateFP8Truncf", "saturate-fp8-truncf", "bool",
+           /*default=*/"false",
+           "Use saturating truncation for 8-bit float types">,
+  ];
 }
 
 //===----------------------------------------------------------------------===//
 // ArithToEmitC
 //===----------------------------------------------------------------------===//
 
-def ArithToEmitCConversionPass : Pass<"convert-arith-to-emitc"> {
-  let summary = "Convert Arith ops to EmitC ops";
-  let description = [{
-    Convert `arith` operations to operations in the `emitc` dialect.
-  }];
+def ConvertArithToEmitC : Pass<"convert-arith-to-emitc"> {
+  let summary = "Convert Arith dialect to EmitC dialect";
   let dependentDialects = ["emitc::EmitCDialect"];
 }
 

mlir/lib/Conversion/ArithToEmitC/ArithToEmitC.cpp

@@ -1,104 +1,60 @@
-//===- ArithToEmitC.cpp - Arith to EmitC conversion -----------------------===//
+//===- ArithToEmitC.cpp - Arith to EmitC Patterns ---------------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
-// This file implements a pass to convert arith ops into emitc ops.
+// This file implements patterns to convert the Arith dialect to the EmitC
+// dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/ArithToEmitC/ArithToEmitC.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/EmitC/IR/EmitC.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/DialectConversion.h"
 
-namespace mlir {
-#define GEN_PASS_DEF_ARITHTOEMITCCONVERSIONPASS
-#include "mlir/Conversion/Passes.h.inc"
-} // namespace mlir
-
 using namespace mlir;
 
-namespace {
-
-static bool isConvertibleToEmitC(Type type) {
-  Type baseType = type;
-  if (auto tensorType = dyn_cast<TensorType>(type)) {
-    if (!tensorType.hasRank() || !tensorType.hasStaticShape()) {
-      return false;
-    }
-    baseType = tensorType.getElementType();
-  }
-
-  if (isa<IndexType>(baseType)) {
-    return true;
-  }
-
-  if (auto intType = dyn_cast<IntegerType>(baseType)) {
-    switch (intType.getWidth()) {
-    case 1:
-    case 8:
-    case 16:
-    case 32:
-    case 64:
-      return true;
-    }
-    return false;
-  }
-
-  if (auto floatType = dyn_cast<FloatType>(baseType)) {
-    return floatType.isF32() || floatType.isF64();
-  }
-
-  return false;
-}
+//===----------------------------------------------------------------------===//
+// Conversion Patterns
+//===----------------------------------------------------------------------===//
 
-class ArithConstantOpConversionPattern
-    : public OpRewritePattern<arith::ConstantOp> {
+namespace {
+template <typename ArithOp, typename EmitCOp>
+class ArithOpConversion final : public OpConversionPattern<ArithOp> {
 public:
-  using OpRewritePattern::OpRewritePattern;
+  using OpConversionPattern<ArithOp>::OpConversionPattern;
 
-  LogicalResult matchAndRewrite(arith::ConstantOp arithConst,
-                                PatternRewriter &rewriter) const override {
+  LogicalResult
+  matchAndRewrite(ArithOp arithOp, typename ArithOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
 
-    auto constantType = arithConst.getType();
-    if (!isConvertibleToEmitC(constantType)) {
-      return rewriter.notifyMatchFailure(arithConst.getLoc(),
-                                         "Type cannot be converted to emitc");
-    }
+    rewriter.template replaceOpWithNewOp<EmitCOp>(arithOp, arithOp.getType(),
+                                                  adaptor.getOperands());
 
-    rewriter.replaceOpWithNewOp<emitc::ConstantOp>(arithConst, constantType,
-                                                   arithConst.getValue());
     return success();
   }
 };
-
-struct ConvertArithToEmitCPass
-    : public impl::ArithToEmitCConversionPassBase<ConvertArithToEmitCPass> {
-public:
-  void runOnOperation() override {
-
-    ConversionTarget target(getContext());
-    target.addIllegalDialect<arith::ArithDialect>();
-    target.addLegalDialect<emitc::EmitCDialect>();
-    RewritePatternSet patterns(&getContext());
-    populateArithToEmitCConversionPatterns(patterns);
-
-    if (failed(applyPartialConversion(getOperation(), target,
-                                      std::move(patterns)))) {
-      signalPassFailure();
-    }
-  }
-};
-
 } // namespace
 
-void mlir::populateArithToEmitCConversionPatterns(RewritePatternSet &patterns) {
-  patterns.add<ArithConstantOpConversionPattern>(patterns.getContext());
+//===----------------------------------------------------------------------===//
+// Pattern population
+//===----------------------------------------------------------------------===//
+
+void mlir::populateArithToEmitCPatterns(TypeConverter &typeConverter,
+                                        RewritePatternSet &patterns) {
+  MLIRContext *ctx = patterns.getContext();
+
+  // clang-format off
+  patterns.add<
+    ArithOpConversion<arith::AddFOp, emitc::AddOp>,
+    ArithOpConversion<arith::DivFOp, emitc::DivOp>,
+    ArithOpConversion<arith::MulFOp, emitc::MulOp>,
+    ArithOpConversion<arith::SubFOp, emitc::SubOp>
+  >(typeConverter, ctx);
+  // clang-format on
 }

mlir/lib/Conversion/ArithToEmitC/ArithToEmitCPass.cpp

@@ -0,0 +1,53 @@
+//===- ArithToEmitCPass.cpp - Arith to EmitC Pass ---------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to convert the Arith dialect to the EmitC
+// dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/ArithToEmitC/ArithToEmitCPass.h"
+
+#include "mlir/Conversion/ArithToEmitC/ArithToEmitC.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/EmitC/IR/EmitC.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_CONVERTARITHTOEMITC
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+struct ConvertArithToEmitC
+    : public impl::ConvertArithToEmitCBase<ConvertArithToEmitC> {
+  void runOnOperation() override;
+};
+} // namespace
+
+void ConvertArithToEmitC::runOnOperation() {
+  ConversionTarget target(getContext());
+
+  target.addLegalDialect<emitc::EmitCDialect>();
+  target.addIllegalDialect<arith::ArithDialect>();
+  target.addLegalOp<arith::ConstantOp>();
+
+  RewritePatternSet patterns(&getContext());
+
+  TypeConverter typeConverter;
+  typeConverter.addConversion([](Type type) { return type; });
+
+  populateArithToEmitCPatterns(typeConverter, patterns);
+
+  if (failed(
+          applyPartialConversion(getOperation(), target, std::move(patterns))))
+    signalPassFailure();
+}

mlir/lib/Conversion/ArithToEmitC/CMakeLists.txt

@@ -1,17 +1,16 @@
-add_mlir_conversion_library(ArithToEmitC
+add_mlir_conversion_library(MLIRArithToEmitC
   ArithToEmitC.cpp
+  ArithToEmitCPass.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/ArithToEmitC
 
   DEPENDS
   MLIRConversionPassIncGen
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
-  MLIREmitCDialect
   MLIRArithDialect
-  MLIRTransforms
-)
+  MLIREmitCDialect
+  MLIRPass
+  MLIRTransformUtils
+  )

mlir/test/Conversion/ArithToEmitC/arith-to-emitc.mlir

@@ -1,21 +1,14 @@
-// RUN: mlir-opt -split-input-file -convert-arith-to-emitc %s | FileCheck %s
+// RUN: mlir-opt -convert-arith-to-emitc %s | FileCheck %s
+
+func.func @arith_ops(%arg0: f32, %arg1: f32) {
+  // CHECK: [[V0:[^ ]*]] = emitc.add %arg0, %arg1 : (f32, f32) -> f32
+  %0 = arith.addf %arg0, %arg1 : f32
+  // CHECK: [[V1:[^ ]*]] = emitc.div %arg0, %arg1 : (f32, f32) -> f32
+  %1 = arith.divf %arg0, %arg1 : f32  
+  // CHECK: [[V2:[^ ]*]] = emitc.mul %arg0, %arg1 : (f32, f32) -> f32
+  %2 = arith.mulf %arg0, %arg1 : f32
+  // CHECK: [[V3:[^ ]*]] = emitc.sub %arg0, %arg1 : (f32, f32) -> f32
+  %3 = arith.subf %arg0, %arg1 : f32
 
-// CHECK-LABEL: arith_constants
-func.func @arith_constants() {
-  // CHECK: emitc.constant
-  // CHECK-SAME: value = 0 : index
-  %c_index = arith.constant 0 : index
-  // CHECK: emitc.constant
-  // CHECK-SAME: value = 0 : i32
-  %c_signless_int_32 = arith.constant 0 : i32
-  // CHECK: emitc.constant
-  // CHECK-SAME: value = 0.{{0+}}e+00 : f32
-  %c_float_32 = arith.constant 0.0 : f32
-  // CHECK: emitc.constant
-  // CHECK-SAME: value = dense<0> : tensor<i32>
-  %c_tensor_single_value = arith.constant dense<0> : tensor<i32>
-  // CHECK: emitc.constant
-  // CHECK-SAME: value{{.*}}[1, 2], [-3, 9], [0, 0], [2, -1]{{.*}}tensor<4x2xi64>
-  %c_tensor_value = arith.constant dense<[[1, 2], [-3, 9], [0, 0], [2, -1]]> : tensor<4x2xi64>
   return
 }

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

@@ -3841,6 +3841,7 @@ cc_library(
         ":AMDGPUToROCDL",
         ":AffineToStandard",
         ":ArithToAMDGPU",
+        ":ArithToEmitC",
         ":ArithToLLVM",
         ":ArithToSPIRV",
         ":ArmNeon2dToIntr",
@@ -7967,6 +7968,32 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "ArithToEmitC",
+    srcs = glob([
+        "lib/Conversion/ArithToEmitC/*.cpp",
+        "lib/Conversion/ArithToEmitC/*.h",
+    ]),
+    hdrs = glob([
+        "include/mlir/Conversion/ArithToEmitC/*.h",
+    ]),
+    includes = [
+        "include",
+        "lib/Conversion/ArithToEmitC",
+    ],
+    deps = [
+        ":ArithDialect",
+        ":ConversionPassIncGen",
+        ":EmitCDialect",
+        ":IR",
+        ":Pass",
+        ":Support",
+        ":TransformUtils",
+        ":Transforms",
+        "//llvm:Support",
+    ],
+)
+
 cc_library(
     name = "ArithToLLVM",
     srcs = glob(["lib/Conversion/ArithToLLVM/*.cpp"]),