[MLIR][SPIRVToLLVM] Conversion patterns for GLSL ops

[georgemitenkov]

Summary

This is the first patch that adds support for GLSL extended
instruction set ops. These are direct conversions, apart from spv.Tan
that is lowered to sin() / cos().

[georgemitenkov]

Added label check to tests and removed unnecessary %{{.*}} = .

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

matchAndRewrite(SPIRVOp operation, ArrayRef<Value> operands,
     }
     Value result = rewriter.template create<LLVMOp>(
         loc, dstType, operation.operand1(), extended);
     rewriter.replaceOp(operation, result);
     return success();
   }
 };
 
+class TanPattern : public SPIRVToLLVMConversion<spirv::GLSLTanOp> {
+public:
+  using SPIRVToLLVMConversion<spirv::GLSLTanOp>::SPIRVToLLVMConversion;
+
+  LogicalResult
+  matchAndRewrite(spirv::GLSLTanOp tanOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto dstType = typeConverter.convertType(tanOp.getType());
+    if (!dstType)
+      return failure();
+
+    Location loc = tanOp.getLoc();
+    Value sin = rewriter.create<LLVM::SinOp>(loc, dstType, tanOp.operand());
+    Value cos = rewriter.create<LLVM::CosOp>(loc, dstType, tanOp.operand());
+    rewriter.replaceOpWithNewOp<LLVM::FDivOp>(tanOp, dstType, sin, cos);
+    return success();
+  }
+};
+
 class VariablePattern : public SPIRVToLLVMConversion<spirv::VariableOp> {
 public:
   using SPIRVToLLVMConversion<spirv::VariableOp>::SPIRVToLLVMConversion;
 
   LogicalResult
   matchAndRewrite(spirv::VariableOp varOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto srcType = varOp.getType();
patterns.insert<
 
       // Control Flow ops
       BranchConversionPattern, BranchConditionalConversionPattern,
       SelectionPattern, MergePattern,
 
       // Function Call op
       FunctionCallPattern,
 
+      // GLSL extended instruction set ops
+      DirectConversionPattern<spirv::GLSLCeilOp, LLVM::FCeilOp>,
+      DirectConversionPattern<spirv::GLSLCosOp, LLVM::CosOp>,
+      DirectConversionPattern<spirv::GLSLExpOp, LLVM::ExpOp>,
+      DirectConversionPattern<spirv::GLSLFAbsOp, LLVM::FAbsOp>,
+      DirectConversionPattern<spirv::GLSLLogOp, LLVM::LogOp>,
+      DirectConversionPattern<spirv::GLSLSinOp, LLVM::SinOp>,
+      DirectConversionPattern<spirv::GLSLSqrtOp, LLVM::SqrtOp>, TanPattern,
+
       // Logical ops
       DirectConversionPattern<spirv::LogicalAndOp, LLVM::AndOp>,
       DirectConversionPattern<spirv::LogicalOrOp, LLVM::OrOp>,
       IComparePattern<spirv::LogicalEqualOp, LLVM::ICmpPredicate::eq>,
       IComparePattern<spirv::LogicalNotEqualOp, LLVM::ICmpPredicate::ne>,
       NotPattern<spirv::LogicalNotOp>,
 
       // Memory ops

mlir/test/Conversion/SPIRVToLLVM/glsl-ops-to-llvm.mlir

+// RUN: mlir-opt -convert-spirv-to-llvm %s | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.Ceil
+//===----------------------------------------------------------------------===//
+
+func @ceil(%arg0: f32, %arg1: vector<3xf16>) {
+  // CHECK: %{{.*}} = "llvm.intr.ceil"(%{{.*}}) : (!llvm.float) -> !llvm.float
+  %0 = spv.GLSL.Ceil %arg0 : f32
+  // CHECK: %{{.*}} = "llvm.intr.ceil"(%{{.*}}) : (!llvm<"<3 x half>">) -> !llvm<"<3 x half>">
+  %1 = spv.GLSL.Ceil %arg1 : vector<3xf16>
+  return
+}
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.Cos
+//===----------------------------------------------------------------------===//
+
+func @cos(%arg0: f32, %arg1: vector<3xf16>) {
+  // CHECK: %{{.*}} = "llvm.intr.cos"(%{{.*}}) : (!llvm.float) -> !llvm.float
+  %0 = spv.GLSL.Cos %arg0 : f32
+  // CHECK: %{{.*}} = "llvm.intr.cos"(%{{.*}}) : (!llvm<"<3 x half>">) -> !llvm<"<3 x half>">
+  %1 = spv.GLSL.Cos %arg1 : vector<3xf16>
+  return
+}
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.Exp
+//===----------------------------------------------------------------------===//
+
+func @exp(%arg0: f32, %arg1: vector<3xf16>) {
+  // CHECK: %{{.*}} = "llvm.intr.exp"(%{{.*}}) : (!llvm.float) -> !llvm.float
+  %0 = spv.GLSL.Exp %arg0 : f32
+  // CHECK: %{{.*}} = "llvm.intr.exp"(%{{.*}}) : (!llvm<"<3 x half>">) -> !llvm<"<3 x half>">
+  %1 = spv.GLSL.Exp %arg1 : vector<3xf16>
+  return
+}
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.FAbs
+//===----------------------------------------------------------------------===//
+
+func @fabs(%arg0: f32, %arg1: vector<3xf16>) {
+  // CHECK: %{{.*}} = "llvm.intr.fabs"(%{{.*}}) : (!llvm.float) -> !llvm.float
+  %0 = spv.GLSL.FAbs %arg0 : f32
+  // CHECK: %{{.*}} = "llvm.intr.fabs"(%{{.*}}) : (!llvm<"<3 x half>">) -> !llvm<"<3 x half>">
+  %1 = spv.GLSL.FAbs %arg1 : vector<3xf16>
+  return
+}
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.Log
+//===----------------------------------------------------------------------===//
+
+func @log(%arg0: f32, %arg1: vector<3xf16>) {
+  // CHECK: %{{.*}} = "llvm.intr.log"(%{{.*}}) : (!llvm.float) -> !llvm.float
+  %0 = spv.GLSL.Log %arg0 : f32
+  // CHECK: %{{.*}} = "llvm.intr.log"(%{{.*}}) : (!llvm<"<3 x half>">) -> !llvm<"<3 x half>">
+  %1 = spv.GLSL.Log %arg1 : vector<3xf16>
+  return
+}
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.Sin
+//===----------------------------------------------------------------------===//
+
+func @sin(%arg0: f32, %arg1: vector<3xf16>) {
+  // CHECK: %{{.*}} = "llvm.intr.sin"(%{{.*}}) : (!llvm.float) -> !llvm.float
+  %0 = spv.GLSL.Sin %arg0 : f32
+  // CHECK: %{{.*}} = "llvm.intr.sin"(%{{.*}}) : (!llvm<"<3 x half>">) -> !llvm<"<3 x half>">
+  %1 = spv.GLSL.Sin %arg1 : vector<3xf16>
+  return
+}
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.Sqrt
+//===----------------------------------------------------------------------===//
+
+func @sqrt(%arg0: f32, %arg1: vector<3xf16>) {
+  // CHECK: %{{.*}} = "llvm.intr.sqrt"(%{{.*}}) : (!llvm.float) -> !llvm.float
+  %0 = spv.GLSL.Sqrt %arg0 : f32
+  // CHECK: %{{.*}} = "llvm.intr.sqrt"(%{{.*}}) : (!llvm<"<3 x half>">) -> !llvm<"<3 x half>">
+  %1 = spv.GLSL.Sqrt %arg1 : vector<3xf16>
+  return
+}
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.Tan
+//===----------------------------------------------------------------------===//
+
+func @tan(%arg0: f32) {
+  // CHECK: %[[SIN:.*]] = "llvm.intr.sin"(%{{.*}}) : (!llvm.float) -> !llvm.float
+  // CHECK: %[[COS:.*]] = "llvm.intr.cos"(%{{.*}}) : (!llvm.float) -> !llvm.float
+  // CHECK: %{{.*}} = llvm.fdiv %[[SIN]], %[[COS]] : !llvm.float
+  %0 = spv.GLSL.Tan %arg0 : f32
+  return
+}