From 108a320a58b13f7adc238d91ab1895ad4569df7a Mon Sep 17 00:00:00 2001
From: Weiwei Li <weiwei.li1@huawei.com>
Date: Fri, 16 Jul 2021 07:25:56 +0800
Subject: [PATCH] [mlir][spirv] Add support for GLSL FMix

Add spv.GLSL.FMix opertaion.

co-authered-by: Alan Liu <alanliu.yf@gmail.com>

Reviewed By: mravishankar

Differential Revision: https://reviews.llvm.org/D104153
---
 mlir/include/mlir/Dialect/SPIRV/IR/SPIRVGLSLOps.td | 38 ++++++++++++++++++++++
 mlir/test/Dialect/SPIRV/IR/glsl-ops.mlir           | 20 ++++++++++++
 mlir/test/Target/SPIRV/glsl-ops.mlir               |  2 ++
 3 files changed, 60 insertions(+)

diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVGLSLOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVGLSLOps.td
index 2097498..4653267 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVGLSLOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVGLSLOps.td
@@ -1123,5 +1123,43 @@ def SPV_GLSLLdexpOp :
   let verifier = [{ return ::verify(*this); }];
 }
 
+def SPV_GLSLFMixOp :
+  SPV_GLSLOp<"FMix", 46, [
+      NoSideEffect, AllTypesMatch<["x", "y", "a", "result"]>]> {
+  let summary = "Builds the linear blend of x and y";
+
+  let description = [{
+    Result is the linear blend of x and y, i.e., x * (1 - a) + y * a.
+
+    The operands must all be a scalar or vector whose component type is floating-point.
+
+    Result Type and the type of all operands must be the same type. Results are computed per component.
+
+    <!-- End of AutoGen section -->
+
+    #### Example:
+
+    ```mlir
+    %0 = spv.GLSL.FMix %x : f32, %y : f32, %a : f32 -> f32
+    %0 = spv.GLSL.FMix %x : vector<4xf32>, %y : vector<4xf32>, %a : vector<4xf32> -> vector<4xf32>
+    ```
+  }];
+
+  let arguments = (ins
+    SPV_ScalarOrVectorOf<SPV_Float>:$x,
+    SPV_ScalarOrVectorOf<SPV_Float>:$y,
+    SPV_ScalarOrVectorOf<SPV_Float>:$a
+  );
+
+  let results = (outs
+    SPV_ScalarOrVectorOf<SPV_Float>:$result
+  );
+
+  let assemblyFormat = [{
+    attr-dict $x `:` type($x) `,` $y `:` type($y) `,` $a `:` type($a) `->` type($result)
+  }];
+
+  let verifier = [{ return success(); }];
+}
 
 #endif // MLIR_DIALECT_SPIRV_IR_GLSL_OPS
diff --git a/mlir/test/Dialect/SPIRV/IR/glsl-ops.mlir b/mlir/test/Dialect/SPIRV/IR/glsl-ops.mlir
index 5b6d382..cb941ad 100644
--- a/mlir/test/Dialect/SPIRV/IR/glsl-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/glsl-ops.mlir
@@ -463,3 +463,23 @@ func @ldexp_wrong_type_vec_2(%arg0 : vector<3xf32>, %arg1 : vector<2xi32>) -> ()
   %0 = spv.GLSL.Ldexp %arg0 : vector<3xf32>, %arg1 : vector<2xi32> -> vector<3xf32>
   return
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GLSL.FMix
+//===----------------------------------------------------------------------===//
+
+func @fmix(%arg0 : f32, %arg1 : f32, %arg2 : f32) -> () {
+  // CHECK: {{%.*}} = spv.GLSL.FMix {{%.*}} : f32, {{%.*}} : f32, {{%.*}} : f32 -> f32
+  %0 = spv.GLSL.FMix %arg0 : f32, %arg1 : f32, %arg2 : f32 -> f32
+  return
+}
+
+// -----
+func @fmix_vector(%arg0 : vector<3xf32>, %arg1 : vector<3xf32>, %arg2 : vector<3xf32>) -> () {
+  // CHECK: {{%.*}} = spv.GLSL.FMix {{%.*}} : vector<3xf32>, {{%.*}} : vector<3xf32>, {{%.*}} : vector<3xf32> -> vector<3xf32>
+  %0 = spv.GLSL.FMix %arg0 : vector<3xf32>, %arg1 : vector<3xf32>, %arg2 : vector<3xf32> -> vector<3xf32>
+  return
+}
+
diff --git a/mlir/test/Target/SPIRV/glsl-ops.mlir b/mlir/test/Target/SPIRV/glsl-ops.mlir
index 8fb5832..1423eed 100644
--- a/mlir/test/Target/SPIRV/glsl-ops.mlir
+++ b/mlir/test/Target/SPIRV/glsl-ops.mlir
@@ -32,6 +32,8 @@ spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
     %13 = spv.GLSL.FrexpStruct %arg0 : f32 -> !spv.struct<(f32, i32)>
     // CHECK: {{%.*}} = spv.GLSL.Ldexp {{%.*}} : f32, {{%.*}} : i32 -> f32
     %14 = spv.GLSL.Ldexp %arg0 : f32, %arg2 : i32 -> f32
+    // CHECK: {{%.*}} = spv.GLSL.FMix {{%.*}} : f32, {{%.*}} : f32, {{%.*}} : f32 -> f32
+    %15 = spv.GLSL.FMix %arg0 : f32, %arg1 : f32, %arg0 : f32 -> f32
     spv.Return
   }
 
-- 
2.7.4