[mlir] Add support for ExpM1 to GLSL/OpenCL SPIRV Backends

[rsuderman]

Summary

Adding a similar decomposition for exponential minus one to the SPIRV
backends along with the necessary tests.

[rsuderman]

Synced to head.

[antiagainst]

Thanks! Just two nits. :)

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

48

I think we can use operation.getType() directly here,

49

Use adaptor.getOperand()?

[rsuderman]

Fixed for nits.

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

 // Operation conversion
 //===----------------------------------------------------------------------===//
 
 // Note that DRR cannot be used for the patterns in this file: we may need to
 // convert type along the way, which requires ConversionPattern. DRR generates
 // normal RewritePattern.
 
 namespace {
+/// Converts math.expm1 to SPIR-V ops.
+///
+/// SPIR-V does not have a direct operations for exp(x)-1. Explicitly lower to
+/// these operations.
+template <typename ExpOp>
+class ExpM1OpPattern final : public OpConversionPattern<math::ExpM1Op> {
+public:
+  using OpConversionPattern<math::ExpM1Op>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(math::ExpM1Op operation, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    assert(adaptor.getOperands().size() == 1);
+    Location loc = operation.getLoc();
+    auto type = this->getTypeConverter()->convertType(operation.getType());
+    auto exp = rewriter.create<ExpOp>(loc, type, adaptor.getOperand());
+    auto one = spirv::ConstantOp::getOne(type, loc, rewriter);
+    rewriter.replaceOpWithNewOp<spirv::FSubOp>(operation, exp, one);
+    return success();
+  }
+};
+
 /// Converts math.log1p to SPIR-V ops.
 ///
 /// SPIR-V does not have a direct operations for log(1+x). Explicitly lower to
 /// these operations.
 template <typename LogOp>
 class Log1pOpPattern final : public OpConversionPattern<math::Log1pOp> {
 public:
   using OpConversionPattern<math::Log1pOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(math::Log1pOp operation, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     assert(adaptor.getOperands().size() == 1);
     Location loc = operation.getLoc();
-    auto type =
-        this->getTypeConverter()->convertType(operation.getOperand().getType());
+    auto type = this->getTypeConverter()->convertType(operation.getType());
     auto one = spirv::ConstantOp::getOne(type, operation.getLoc(), rewriter);
     auto onePlus =
-        rewriter.create<spirv::FAddOp>(loc, one, adaptor.getOperands()[0]);
+        rewriter.create<spirv::FAddOp>(loc, one, adaptor.getOperand());
     rewriter.replaceOpWithNewOp<LogOp>(operation, type, onePlus);
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Pattern population
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 void populateMathToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                  RewritePatternSet &patterns) {
 
   // GLSL patterns
   patterns
-      .add<Log1pOpPattern<spirv::GLSLLogOp>,
+      .add<Log1pOpPattern<spirv::GLSLLogOp>, ExpM1OpPattern<spirv::GLSLExpOp>,
            spirv::ElementwiseOpPattern<math::AbsOp, spirv::GLSLFAbsOp>,
            spirv::ElementwiseOpPattern<math::CeilOp, spirv::GLSLCeilOp>,
            spirv::ElementwiseOpPattern<math::CosOp, spirv::GLSLCosOp>,
            spirv::ElementwiseOpPattern<math::ExpOp, spirv::GLSLExpOp>,
            spirv::ElementwiseOpPattern<math::FloorOp, spirv::GLSLFloorOp>,
            spirv::ElementwiseOpPattern<math::LogOp, spirv::GLSLLogOp>,
            spirv::ElementwiseOpPattern<math::PowFOp, spirv::GLSLPowOp>,
            spirv::ElementwiseOpPattern<math::RsqrtOp, spirv::GLSLInverseSqrtOp>,
            spirv::ElementwiseOpPattern<math::SinOp, spirv::GLSLSinOp>,
            spirv::ElementwiseOpPattern<math::SqrtOp, spirv::GLSLSqrtOp>,
            spirv::ElementwiseOpPattern<math::TanhOp, spirv::GLSLTanhOp>,
            spirv::ElementwiseOpPattern<math::FmaOp, spirv::GLSLFmaOp>>(
           typeConverter, patterns.getContext());
 
   // OpenCL patterns
-  patterns.add<Log1pOpPattern<spirv::OCLLogOp>,
+  patterns.add<Log1pOpPattern<spirv::OCLLogOp>, ExpM1OpPattern<spirv::OCLExpOp>,
                spirv::ElementwiseOpPattern<math::AbsOp, spirv::OCLFAbsOp>,
                spirv::ElementwiseOpPattern<math::CeilOp, spirv::OCLCeilOp>,
                spirv::ElementwiseOpPattern<math::CosOp, spirv::OCLCosOp>,
                spirv::ElementwiseOpPattern<math::ErfOp, spirv::OCLErfOp>,
                spirv::ElementwiseOpPattern<math::ExpOp, spirv::OCLExpOp>,
                spirv::ElementwiseOpPattern<math::FloorOp, spirv::OCLFloorOp>,
                spirv::ElementwiseOpPattern<math::LogOp, spirv::OCLLogOp>,
                spirv::ElementwiseOpPattern<math::PowFOp, spirv::OCLPowOp>,
                spirv::ElementwiseOpPattern<math::RsqrtOp, spirv::OCLRsqrtOp>,
                spirv::ElementwiseOpPattern<math::SinOp, spirv::OCLSinOp>,
                spirv::ElementwiseOpPattern<math::SqrtOp, spirv::OCLSqrtOp>,
                spirv::ElementwiseOpPattern<math::TanhOp, spirv::OCLTanhOp>>(
       typeConverter, patterns.getContext());
 }
 
 } // namespace mlir

mlir/test/Conversion/MathToSPIRV/math-to-glsl-spirv.mlir

 // RUN: mlir-opt -split-input-file -convert-math-to-spirv -verify-diagnostics %s -o - | FileCheck %s
 
 module attributes { spv.target_env = #spv.target_env<#spv.vce<v1.0, [Shader], []>, {}> } {
 
 // CHECK-LABEL: @float32_unary_scalar
 func @float32_unary_scalar(%arg0: f32) {
   // CHECK: spv.GLSL.Cos %{{.*}}: f32
   %0 = math.cos %arg0 : f32
   // CHECK: spv.GLSL.Exp %{{.*}}: f32
   %1 = math.exp %arg0 : f32
+  // CHECK: %[[EXP:.+]] = spv.GLSL.Exp %arg0
+  // CHECK: %[[ONE:.+]] = spv.Constant 1.000000e+00 : f32
+  // CHECK: spv.FSub %[[EXP]], %[[ONE]]
+  %2 = math.expm1 %arg0 : f32
   // CHECK: spv.GLSL.Log %{{.*}}: f32
-  %2 = math.log %arg0 : f32
+  %3 = math.log %arg0 : f32
   // CHECK: %[[ONE:.+]] = spv.Constant 1.000000e+00 : f32
   // CHECK: %[[ADDONE:.+]] = spv.FAdd %[[ONE]], %{{.+}}
   // CHECK: spv.GLSL.Log %[[ADDONE]]
-  %3 = math.log1p %arg0 : f32
+  %4 = math.log1p %arg0 : f32
   // CHECK: spv.GLSL.InverseSqrt %{{.*}}: f32
-  %4 = math.rsqrt %arg0 : f32
+  %5 = math.rsqrt %arg0 : f32
   // CHECK: spv.GLSL.Sqrt %{{.*}}: f32
-  %5 = math.sqrt %arg0 : f32
+  %6 = math.sqrt %arg0 : f32
   // CHECK: spv.GLSL.Tanh %{{.*}}: f32
-  %6 = math.tanh %arg0 : f32
+  %7 = math.tanh %arg0 : f32
   // CHECK: spv.GLSL.Sin %{{.*}}: f32
-  %7 = math.sin %arg0 : f32
+  %8 = math.sin %arg0 : f32
   // CHECK: spv.GLSL.FAbs %{{.*}}: f32
-  %8 = math.abs %arg0 : f32
+  %9 = math.abs %arg0 : f32
   // CHECK: spv.GLSL.Ceil %{{.*}}: f32
-  %9 = math.ceil %arg0 : f32
+  %10 = math.ceil %arg0 : f32
   // CHECK: spv.GLSL.Floor %{{.*}}: f32
-  %10 = math.floor %arg0 : f32
+  %11 = math.floor %arg0 : f32
   return
 }
 
 // CHECK-LABEL: @float32_unary_vector
 func @float32_unary_vector(%arg0: vector<3xf32>) {
   // CHECK: spv.GLSL.Cos %{{.*}}: vector<3xf32>
   %0 = math.cos %arg0 : vector<3xf32>
   // CHECK: spv.GLSL.Exp %{{.*}}: vector<3xf32>
   %1 = math.exp %arg0 : vector<3xf32>
+  // CHECK: %[[EXP:.+]] = spv.GLSL.Exp %arg0
+  // CHECK: %[[ONE:.+]] = spv.Constant dense<1.000000e+00> : vector<3xf32>
+  // CHECK: spv.FSub %[[EXP]], %[[ONE]]
+  %2 = math.expm1 %arg0 : vector<3xf32>
   // CHECK: spv.GLSL.Log %{{.*}}: vector<3xf32>
-  %2 = math.log %arg0 : vector<3xf32>
+  %3 = math.log %arg0 : vector<3xf32>
   // CHECK: %[[ONE:.+]] = spv.Constant dense<1.000000e+00> : vector<3xf32>
   // CHECK: %[[ADDONE:.+]] = spv.FAdd %[[ONE]], %{{.+}}
   // CHECK: spv.GLSL.Log %[[ADDONE]]
-  %3 = math.log1p %arg0 : vector<3xf32>
+  %4 = math.log1p %arg0 : vector<3xf32>
   // CHECK: spv.GLSL.InverseSqrt %{{.*}}: vector<3xf32>
-  %4 = math.rsqrt %arg0 : vector<3xf32>
+  %5 = math.rsqrt %arg0 : vector<3xf32>
   // CHECK: spv.GLSL.Sqrt %{{.*}}: vector<3xf32>
-  %5 = math.sqrt %arg0 : vector<3xf32>
+  %6 = math.sqrt %arg0 : vector<3xf32>
   // CHECK: spv.GLSL.Tanh %{{.*}}: vector<3xf32>
-  %6 = math.tanh %arg0 : vector<3xf32>
+  %7 = math.tanh %arg0 : vector<3xf32>
   // CHECK: spv.GLSL.Sin %{{.*}}: vector<3xf32>
-  %7 = math.sin %arg0 : vector<3xf32>
+  %8 = math.sin %arg0 : vector<3xf32>
   return
 }
 
 // CHECK-LABEL: @float32_binary_scalar
 func @float32_binary_scalar(%lhs: f32, %rhs: f32) {
   // CHECK: spv.GLSL.Pow %{{.*}}: f32
   %0 = math.powf %lhs, %rhs : f32
   return

mlir/test/Conversion/MathToSPIRV/math-to-opencl-spirv.mlir

 // RUN: mlir-opt -split-input-file -convert-math-to-spirv -verify-diagnostics %s -o - | FileCheck %s
 
 module attributes { spv.target_env = #spv.target_env<#spv.vce<v1.0, [Kernel], []>, {}> } {
 
 // CHECK-LABEL: @float32_unary_scalar
 func @float32_unary_scalar(%arg0: f32) {
   // CHECK: spv.OCL.cos %{{.*}}: f32
   %0 = math.cos %arg0 : f32
   // CHECK: spv.OCL.exp %{{.*}}: f32
   %1 = math.exp %arg0 : f32
+  // CHECK: %[[EXP:.+]] = spv.OCL.exp %arg0
+  // CHECK: %[[ONE:.+]] = spv.Constant 1.000000e+00 : f32
+  // CHECK: spv.FSub %[[EXP]], %[[ONE]]
+  %2 = math.expm1 %arg0 : f32
   // CHECK: spv.OCL.log %{{.*}}: f32
-  %2 = math.log %arg0 : f32
+  %3 = math.log %arg0 : f32
   // CHECK: %[[ONE:.+]] = spv.Constant 1.000000e+00 : f32
   // CHECK: %[[ADDONE:.+]] = spv.FAdd %[[ONE]], %{{.+}}
   // CHECK: spv.OCL.log %[[ADDONE]]
-  %3 = math.log1p %arg0 : f32
+  %4 = math.log1p %arg0 : f32
   // CHECK: spv.OCL.rsqrt %{{.*}}: f32
-  %4 = math.rsqrt %arg0 : f32
+  %5 = math.rsqrt %arg0 : f32
   // CHECK: spv.OCL.sqrt %{{.*}}: f32
-  %5 = math.sqrt %arg0 : f32
+  %6 = math.sqrt %arg0 : f32
   // CHECK: spv.OCL.tanh %{{.*}}: f32
-  %6 = math.tanh %arg0 : f32
+  %7 = math.tanh %arg0 : f32
   // CHECK: spv.OCL.sin %{{.*}}: f32
-  %7 = math.sin %arg0 : f32
+  %8 = math.sin %arg0 : f32
   // CHECK: spv.OCL.fabs %{{.*}}: f32
-  %8 = math.abs %arg0 : f32
+  %9 = math.abs %arg0 : f32
   // CHECK: spv.OCL.ceil %{{.*}}: f32
-  %9 = math.ceil %arg0 : f32
+  %10 = math.ceil %arg0 : f32
   // CHECK: spv.OCL.floor %{{.*}}: f32
-  %10 = math.floor %arg0 : f32
+  %11 = math.floor %arg0 : f32
   // CHECK: spv.OCL.erf %{{.*}}: f32
-  %11 = math.erf %arg0 : f32
+  %12 = math.erf %arg0 : f32
   return
 }
 
 // CHECK-LABEL: @float32_unary_vector
 func @float32_unary_vector(%arg0: vector<3xf32>) {
   // CHECK: spv.OCL.cos %{{.*}}: vector<3xf32>
   %0 = math.cos %arg0 : vector<3xf32>
   // CHECK: spv.OCL.exp %{{.*}}: vector<3xf32>
   %1 = math.exp %arg0 : vector<3xf32>
+  // CHECK: %[[EXP:.+]] = spv.OCL.exp %arg0
+  // CHECK: %[[ONE:.+]] = spv.Constant dense<1.000000e+00> : vector<3xf32>
+  // CHECK: spv.FSub %[[EXP]], %[[ONE]]
+  %2 = math.expm1 %arg0 : vector<3xf32>
   // CHECK: spv.OCL.log %{{.*}}: vector<3xf32>
-  %2 = math.log %arg0 : vector<3xf32>
+  %3 = math.log %arg0 : vector<3xf32>
   // CHECK: %[[ONE:.+]] = spv.Constant dense<1.000000e+00> : vector<3xf32>
   // CHECK: %[[ADDONE:.+]] = spv.FAdd %[[ONE]], %{{.+}}
   // CHECK: spv.OCL.log %[[ADDONE]]
-  %3 = math.log1p %arg0 : vector<3xf32>
+  %4 = math.log1p %arg0 : vector<3xf32>
   // CHECK: spv.OCL.rsqrt %{{.*}}: vector<3xf32>
-  %4 = math.rsqrt %arg0 : vector<3xf32>
+  %5 = math.rsqrt %arg0 : vector<3xf32>
   // CHECK: spv.OCL.sqrt %{{.*}}: vector<3xf32>
-  %5 = math.sqrt %arg0 : vector<3xf32>
+  %6 = math.sqrt %arg0 : vector<3xf32>
   // CHECK: spv.OCL.tanh %{{.*}}: vector<3xf32>
-  %6 = math.tanh %arg0 : vector<3xf32>
+  %7 = math.tanh %arg0 : vector<3xf32>
   // CHECK: spv.OCL.sin %{{.*}}: vector<3xf32>
-  %7 = math.sin %arg0 : vector<3xf32>
+  %8 = math.sin %arg0 : vector<3xf32>
   return
 }
 
 // CHECK-LABEL: @float32_binary_scalar
 func @float32_binary_scalar(%lhs: f32, %rhs: f32) {
   // CHECK: spv.OCL.pow %{{.*}}: f32
   %0 = math.powf %lhs, %rhs : f32
   return