[mlir] [LLVMIR] add all vector reduction intrinsics to LLVM IR dialect

[aartbik]

Summary

This allows for lowering of VectorOps (and others) into a LLVM IR
that maps directly to efficient implementations on the target machines.

http://llvm.org/docs/LangRef.html#experimental-vector-reduction-intrinsics

[ftynse]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

796

Something went wrong here with braces, mismerge?

[nicolasvasilache]

mlir/test/Target/llvmir-intrinsics.mlir

102

Did you try to pipe these through mlir-translate --mlir-to-llvmir ?
I've had surprises on my end with more complex masked operations.

[nicolasvasilache]

mlir/test/Target/llvmir-intrinsics.mlir

102

I would also argue that we should have an extra test in this file that does the piping through mlir-translate --mlir-to-llvmir (no need to filecheck the result just that it doesn't crash).
However @ftynse repeatedly made the case against it.
But I don't buy the separability of tests argument here :) esp in light of how easy it is to get crashes when trying to go all the way down to actual LLVMIR.

[rriddle]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

790

Can you define a template for these? The builder is effectively the same for all of these.

[aartbik]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

796

I took the code as generated by your awesome utility, and only reformatted on the 80-col violation. This bracket was placed here.

Is there any way to format a *.td file automatically?

mlir/test/Target/llvmir-intrinsics.mlir

102

Yes Sir! I typically tests my examples with mlir-opt and mlir-translate to be sure. They only time I saw a crash was using an integer intrinsic on a floating-point value by mistake (something we may want to more gracefully detect at some point).

102

Ack on the test, but I am guessing not in this CL ;-)

[aartbik]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

790

Ha! I actually started to write a template by hand and then was pointed to Alex' brilliant tool to generate this code automatically. I am okay cleaning up the code by hand now, but then some of the automation is lost :-)

[ftynse]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

790

We should eventually injecting the generated code into the file instead of copy-pasting, but for now it's fine if it stays as is.

796

It's an open project to format those :)

mlir/test/Target/llvmir-intrinsics.mlir

102

@nicolasvasilache this test does exactly mlir-translate --mlir-to-llvmir

[rriddle]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

790

So the plan is to have a GeneratedOps.td alongside this one? That makes sense if so, but can we start a new file for generated ops prematurely to avoid extra work later on? It also keeps this file clean so that people don't look at it as just technical debt laying around to be cleaned up.

[ftynse]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

790

Good idea, let's have a separate file now and clearly mark it as autogenerated to avoid random edits. The generator should already produce a header explaining that.

[aartbik]

less code is better

[aartbik]

Okay, friends of concise code, PTAL

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

 class LLVM_EnumAttr<string name, string llvmName, string description,
                     list<LLVM_EnumAttrCase> cases> :
     I64EnumAttr<name, description, cases> {
 
   // The equivalent enum class name in LLVM.
   string llvmClassName = llvmName;
 }
 
+// LLVM vector reduction over a single vector.
+class LLVM_VectorReduction<string mnem> :
+  LLVM_IntrOp<"experimental.vector.reduce." # mnem, []>,
+  Arguments<(ins LLVM_Type)>, Results<(outs LLVM_Type:$res)> {
+    let llvmBuilder = [{
+      llvm::Module *module = builder.GetInsertBlock()->getModule();
+      llvm::Function *fn = llvm::Intrinsic::getDeclaration(
+        module,
+        llvm::Intrinsic::experimental_vector_reduce_}] #
+        !subst(".","_", mnem) # [{, {
+          opInst.getOperand(0).getType().cast<LLVM::LLVMType>()
+          .getUnderlyingType(),
+        });
+      auto operands = lookupValues(opInst.getOperands());
+      $res = builder.CreateCall(fn, operands);
+    }];
+}
+
+// LLVM vector reduction over a single vector, with an initial value.
+class LLVM_VectorReductionV2<string mnem> :
+  LLVM_IntrOp<"experimental.vector.reduce.v2." # mnem, []>,
+  Arguments<(ins LLVM_Type, LLVM_Type)>, Results<(outs LLVM_Type:$res)> {
+    let llvmBuilder = [{
+      llvm::Module *module = builder.GetInsertBlock()->getModule();
+      llvm::Function *fn = llvm::Intrinsic::getDeclaration(
+        module,
+        llvm::Intrinsic::experimental_vector_reduce_v2_}] #
+        !subst(".","_", mnem) # [{, {
+          opInst.getResult(0).getType().cast<LLVM::LLVMType>()
+          .getUnderlyingType(),
+          opInst.getOperand(1).getType().cast<LLVM::LLVMType>()
+          .getUnderlyingType(),
+        });
+      auto operands = lookupValues(opInst.getOperands());
+      $res = builder.CreateCall(fn, operands);
+    }];
+}
+
 #endif  // LLVMIR_OP_BASE

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

def LLVM_Prefetch : LLVM_ZeroResultOp<"intr.prefetch">,
   let llvmBuilder = [{
     llvm::Module *module = builder.GetInsertBlock()->getModule();
     llvm::Function *fn = llvm::Intrinsic::getDeclaration(
         module, llvm::Intrinsic::prefetch, $addr->getType());
     builder.CreateCall(fn, {$addr, $rw, $hint, $cache});
   }];
 }
 
+//
+// Vector Reductions.
+//
+
+def LLVM_experimental_vector_reduce_add : LLVM_VectorReduction<"add">;
+def LLVM_experimental_vector_reduce_and : LLVM_VectorReduction<"and">;
+def LLVM_experimental_vector_reduce_mul : LLVM_VectorReduction<"mul">;
+def LLVM_experimental_vector_reduce_fmax : LLVM_VectorReduction<"fmax">;
+def LLVM_experimental_vector_reduce_fmin : LLVM_VectorReduction<"fmin">;
+def LLVM_experimental_vector_reduce_or : LLVM_VectorReduction<"or">;
+def LLVM_experimental_vector_reduce_smax : LLVM_VectorReduction<"smax">;
+def LLVM_experimental_vector_reduce_smin : LLVM_VectorReduction<"smin">;
+def LLVM_experimental_vector_reduce_umax : LLVM_VectorReduction<"umax">;
+def LLVM_experimental_vector_reduce_umin : LLVM_VectorReduction<"umin">;
+def LLVM_experimental_vector_reduce_xor : LLVM_VectorReduction<"xor">;
+
+def LLVM_experimental_vector_reduce_v2_fadd : LLVM_VectorReductionV2<"fadd">;
+def LLVM_experimental_vector_reduce_v2_fmul : LLVM_VectorReductionV2<"fmul">;
+
+//
+// Atomic operations.
+//
+
 def AtomicBinOpXchg : I64EnumAttrCase<"xchg", 0>;
 def AtomicBinOpAdd  : I64EnumAttrCase<"add", 1>;
 def AtomicBinOpSub  : I64EnumAttrCase<"sub", 2>;
 def AtomicBinOpAnd  : I64EnumAttrCase<"_and", 3>;
 def AtomicBinOpNand : I64EnumAttrCase<"nand", 4>;
 def AtomicBinOpOr   : I64EnumAttrCase<"_or", 5>;
 def AtomicBinOpXor  : I64EnumAttrCase<"_xor", 6>;
 def AtomicBinOpMax  : I64EnumAttrCase<"max", 7>;

mlir/test/Target/llvmir-intrinsics.mlir

 llvm.func @copysign_test(%arg0: !llvm.float, %arg1: !llvm.float, %arg2: !llvm<"<8 x float>">, %arg3: !llvm<"<8 x float>">) {
   // CHECK: call float @llvm.copysign.f32
   "llvm.intr.copysign"(%arg0, %arg1) : (!llvm.float, !llvm.float) -> !llvm.float
   // CHECK: call <8 x float> @llvm.copysign.v8f32
   "llvm.intr.copysign"(%arg2, %arg3) : (!llvm<"<8 x float>">, !llvm<"<8 x float>">) -> !llvm<"<8 x float>">
   llvm.return
 }
 
+// CHECK-LABEL: @vector_reductions
+llvm.func @vector_reductions(%arg0: !llvm.float, %arg1: !llvm<"<8 x float>">, %arg2: !llvm<"<8 x i32>">) {
+  // CHECK: call i32 @llvm.experimental.vector.reduce.add.v8i32
+  "llvm.intr.experimental.vector.reduce.add"(%arg2) : (!llvm<"<8 x i32>">) -> !llvm.i32
+  // CHECK: call i32 @llvm.experimental.vector.reduce.and.v8i32
+  "llvm.intr.experimental.vector.reduce.and"(%arg2) : (!llvm<"<8 x i32>">) -> !llvm.i32
+  // CHECK: call float @llvm.experimental.vector.reduce.fmax.v8f32
+  "llvm.intr.experimental.vector.reduce.fmax"(%arg1) : (!llvm<"<8 x float>">) -> !llvm.float
+  // CHECK: call float @llvm.experimental.vector.reduce.fmin.v8f32
+  "llvm.intr.experimental.vector.reduce.fmin"(%arg1) : (!llvm<"<8 x float>">) -> !llvm.float
+  // CHECK: call i32 @llvm.experimental.vector.reduce.mul.v8i32
+  "llvm.intr.experimental.vector.reduce.mul"(%arg2) : (!llvm<"<8 x i32>">) -> !llvm.i32
+  // CHECK: call i32 @llvm.experimental.vector.reduce.or.v8i32
+  "llvm.intr.experimental.vector.reduce.or"(%arg2) : (!llvm<"<8 x i32>">) -> !llvm.i32
+  // CHECK: call i32 @llvm.experimental.vector.reduce.smax.v8i32
+  "llvm.intr.experimental.vector.reduce.smax"(%arg2) : (!llvm<"<8 x i32>">) -> !llvm.i32
+  // CHECK: call i32 @llvm.experimental.vector.reduce.smin.v8i32
+  "llvm.intr.experimental.vector.reduce.smin"(%arg2) : (!llvm<"<8 x i32>">) -> !llvm.i32
+  // CHECK: call i32 @llvm.experimental.vector.reduce.umax.v8i32
+  "llvm.intr.experimental.vector.reduce.umax"(%arg2) : (!llvm<"<8 x i32>">) -> !llvm.i32
+  // CHECK: call i32 @llvm.experimental.vector.reduce.umin.v8i32
+  "llvm.intr.experimental.vector.reduce.umin"(%arg2) : (!llvm<"<8 x i32>">) -> !llvm.i32
+  // CHECK: call float @llvm.experimental.vector.reduce.v2.fadd.f32.v8f32
+  "llvm.intr.experimental.vector.reduce.v2.fadd"(%arg0, %arg1) : (!llvm.float, !llvm<"<8 x float>">) -> !llvm.float
+  // CHECK: call float @llvm.experimental.vector.reduce.v2.fmul.f32.v8f32
+  "llvm.intr.experimental.vector.reduce.v2.fmul"(%arg0, %arg1) : (!llvm.float, !llvm<"<8 x float>">) -> !llvm.float
+  // CHECK: call i32 @llvm.experimental.vector.reduce.xor.v8i32
+  "llvm.intr.experimental.vector.reduce.xor"(%arg2) : (!llvm<"<8 x i32>">) -> !llvm.i32
+  llvm.return
+}
+
 // Check that intrinsics are declared with appropriate types.
 // CHECK-DAG: declare float @llvm.fma.f32(float, float, float)
 // CHECK-DAG: declare <8 x float> @llvm.fma.v8f32(<8 x float>, <8 x float>, <8 x float>) #0
 // CHECK-DAG: declare float @llvm.fmuladd.f32(float, float, float)
 // CHECK-DAG: declare <8 x float> @llvm.fmuladd.v8f32(<8 x float>, <8 x float>, <8 x float>) #0
 // CHECK-DAG: declare void @llvm.prefetch.p0i8(i8* nocapture readonly, i32 immarg, i32 immarg, i32)
 // CHECK-DAG: declare float @llvm.exp.f32(float)
 // CHECK-DAG: declare <8 x float> @llvm.exp.v8f32(<8 x float>) #0