Add Math dialect to EmitC conversion support (#729)

- Add MathToEmitC pass and ArithExpandOps pass to StablehloToExecutable
- Add Math Log conversion pattern for f32 and f64 types
- Add cmath include for log/logf functions
- Add test for Math Log to EmitC conversion

GitOrigin-RevId: a6a7f80125e3b92981ca40194181f7ff923d09d2

Author: christopherbate

mlir-tensorrt/compiler/CMakeLists.txt

@@ -39,14 +39,14 @@ mtrt_add_project_targets(MLIRTensorRT
   MLIRAsyncDialect
   MLIRBufferizationTransforms
   MLIRControlFlowTransforms
+  MLIRDebug
   MLIRFuncInlinerExtension
   MLIRNVVMTarget
   MLIRPtrDialect
+  MLIRTargetCpp
   MLIRTargetLLVM
   MLIRTensorInferTypeOpInterfaceImpl
   MLIRTensorTransformOps
-  MLIRTargetCpp
-  MLIRDebug
   MLIRTransformDialectTransforms
   )
 

mlir-tensorrt/compiler/lib/Compiler/StablehloToExecutable/CMakeLists.txt

@@ -41,8 +41,10 @@ add_mlir_tensorrt_library(MLIRTensorRTCompilerStableHloToExecutable
       MLIRTensorRTTransformsUnrollForLoops
       StablehloLinalgTransforms
     MLIR_LIBS PUBLIC
+      MLIRArithTransforms
       MLIREmitCTransforms
       MLIRIR
       MLIRLLVMIRTransforms
+      MLIRMathToEmitC
     )
 

mlir-tensorrt/compiler/lib/Compiler/StablehloToExecutable/StablehloToExecutable.cpp

@@ -37,8 +37,11 @@
 #include "mlir-tensorrt/Transforms/Passes.h"
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/ComplexToStandard/ComplexToStandard.h"
+#include "mlir/Conversion/MathToEmitC/MathToEmitC.h"
+#include "mlir/Conversion/MathToEmitC/MathToEmitCPass.h"
 #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
 #include "mlir/Dialect/Affine/Passes.h"
+#include "mlir/Dialect/Arith/Transforms/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/EmitC/Transforms/Passes.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -241,6 +244,10 @@ void StablehloToExecutableTask::populatePassManager() {
     // For EmitC, just run Host-to-EmitC followed
     // by cleanup and expression forming.
     if (hostTarget == HostTarget::EmitC) {
+      pm.addNestedPass<func::FuncOp>(mlir::createConvertMathToEmitC());
+      pm.addNestedPass<func::FuncOp>(arith::createArithExpandOpsPass());
+      pm.addNestedPass<func::FuncOp>(createCSEPass());
+      pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
       pm.addPass(createConvertHostToEmitCPass({options.artifactsDirectory}));
       addCleanupPasses(pm);
       // The EmitC "form-expressions" pass combines operations into

mlir-tensorrt/compiler/lib/Conversion/HostToEmitC/CMakeLists.txt

@@ -16,4 +16,5 @@ add_mlir_tensorrt_library(MLIRTensorRTHostToEmitC
     MLIRSCFToEmitC
     MLIRFuncToEmitC
     MLIRFuncTransforms
+    MLIRMathDialect
   )

mlir-tensorrt/compiler/lib/Conversion/HostToEmitC/HostToEmitC.cpp

@@ -34,6 +34,7 @@
 #include "mlir/Dialect/EmitC/IR/EmitC.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Func/Transforms/OneToNFuncConversions.h"
+#include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
@@ -1686,7 +1687,33 @@ class MemRefExtractAlignedPointerAsIndexConverter
     return success();
   }
 };
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Math conversions missing from 'math-to-emitc' pass
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct MathLogToEmitCPattern : public EmitCConversionPattern<math::LogOp> {
+  using EmitCConversionPattern::EmitCConversionPattern;
 
+  LogicalResult
+  matchAndRewrite(math::LogOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    Value input = adaptor.getOperand();
+    Type inputType = input.getType();
+    if (!inputType.isF32() && !inputType.isF64())
+      return failure();
+    llvm::StringRef funcName = "log";
+    if (inputType.isF32())
+      funcName = "logf";
+    auto callOp =
+        createCallOpaque(rewriter, op.getLoc(), inputType, funcName, {input});
+    rewriter.replaceOp(op, callOp.getResult(0));
+    return success();
+  }
+};
 } // namespace
 
 //===----------------------------------------------------------------------===//
@@ -1762,6 +1789,7 @@ static void populateEmitCConversionPatternsAndLegality(
       CUDAStreamSyncConverter,
       ExecutorPrintConverter,
       ExtractStridedMetadataOpLowering,
+      MathLogToEmitCPattern,
       MemRefAllocOpLowering,
       MemrefCastOpLowering,
       MemRefDimOpLowering,
@@ -1858,6 +1886,7 @@ class HostToEmitCPass
     rewriter.create<emitc::IncludeOp>(moduleOp->getLoc(), "cstdio", true);
     rewriter.create<emitc::IncludeOp>(moduleOp->getLoc(), "cstdint", true);
     rewriter.create<emitc::IncludeOp>(moduleOp->getLoc(), "cstdlib", true);
+    rewriter.create<emitc::IncludeOp>(moduleOp->getLoc(), "cmath", true);
     rewriter.create<emitc::IncludeOp>(moduleOp->getLoc(), "MTRTRuntime.h",
                                       false);
 

mlir-tensorrt/compiler/test/Conversion/HostToEmitC/math-to-emitc.mlir

@@ -0,0 +1,29 @@
+// RUN: rm -rf %t || true
+// RUN: mkdir -p %t
+// RUN: mlir-tensorrt-opt -split-input-file -convert-host-to-emitc="artifacts-dir=%t" %s | \
+// RUN: mlir-tensorrt-translate -split-input-file -mlir-to-cpp | FileCheck %s --check-prefix=CPP
+
+// Test Math Log conversion to EmitC
+
+func.func @math_log_f32(%arg0: f32) -> f32 {
+  %0 = math.log %arg0 : f32
+  return %0 : f32
+}
+
+// CPP-LABEL: float math_log_f32(float v1) {
+// CPP-NEXT:   float v2 = logf(v1);
+// CPP-NEXT:   return v2;
+// CPP-NEXT: }
+
+// -----
+
+func.func @math_log_f64(%arg0: f64) -> f64 {
+  %0 = math.log %arg0 : f64
+  return %0 : f64
+}
+
+// CPP-LABEL: double math_log_f64(double v1) {
+// CPP-NEXT:   double v2 = log(v1);
+// CPP-NEXT:   return v2;
+// CPP-NEXT: }
+