[mlir][llvm] Import matrix, vector, and assume intrinsics from LLVM IR.

[gysit]

Summary

The revision adds support to import:

• matrix intrinsics
• vector reduce fadd/fmul intrinsics
• assume intrinsics

from LLVM IR.

[ftynse]

While you are looking at these, maybe we can also have more specific type constraints than the generic LLVM_Type?

[gysit]

While you are looking at these, maybe we can also have more specific type constraints than the generic LLVM_Type?

Yes that may actually unveil some import bugs. I'll put it on my TODO list to make the types of at least some of the operation types more precise. It seems like the type definitions are already there!

[ftynse]

Yes, I switched most of the operations to have better type constraints, but didn't do this for intrinsics.

[gysit]

Remove redunant variable name in assume intrinsic declaration.

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

def LLVM_SMulWithOverflowOp
   let arguments = (ins LLVM_Type, LLVM_Type);
 }
 def LLVM_UMulWithOverflowOp
     : LLVM_IntrOp<"umul.with.overflow", [0], [], [], 2> {
   let arguments = (ins LLVM_Type, LLVM_Type);
 }
 
 
-def LLVM_AssumeOp : LLVM_Op<"intr.assume", []> {
+def LLVM_AssumeOp : LLVM_ZeroResultIntrOp<"assume", []> {
   let arguments = (ins LLVM_Type:$cond);
-  let llvmBuilder = [{
-    llvm::Module *module = builder.GetInsertBlock()->getModule();
-    llvm::Function *fn =
-        llvm::Intrinsic::getDeclaration(module, llvm::Intrinsic::assume, {});
-    builder.CreateCall(fn, {$cond});
-  }];
 }
 
 
 //
 // Coroutine intrinsics.
 //
 
 def LLVM_CoroIdOp : LLVM_IntrOp<"coro.id", [], [], [], 1> {

 
 /// Create a column major, strided 2-D matrix load, as specified in the LLVM
 /// MatrixBuilder.
 /// data       - Start address of the matrix read
 /// rows       - Number of rows in matrix (must be a constant)
 /// isVolatile - True if the load operation is marked as volatile.
 /// columns    - Number of columns in matrix (must be a constant)
 /// stride     - Space between columns
-def LLVM_MatrixColumnMajorLoadOp : LLVM_Op<"intr.matrix.column.major.load"> {
+def LLVM_MatrixColumnMajorLoadOp : LLVM_OneResultIntrOp<"matrix.column.major.load"> {
   let arguments = (ins LLVM_Type:$data, LLVM_Type:$stride, I1Attr:$isVolatile,
                    I32Attr:$rows, I32Attr:$columns);
   let results = (outs LLVM_AnyVector:$res);
   let builders = [LLVM_OneResultOpBuilder];
+  let assemblyFormat = "$data `,` `<` `stride` `=` $stride `>` attr-dict"
+    "`:` type($res) `from` type($data) `stride` type($stride)";
+
   string llvmBuilder = [{
     llvm::MatrixBuilder mb(builder);
     const llvm::DataLayout &dl =
       builder.GetInsertBlock()->getModule()->getDataLayout();
     llvm::Type *ElemTy = moduleTranslation.convertType(
         getVectorElementType(op.getType()));
     llvm::Align align = dl.getABITypeAlign(ElemTy);
     $res = mb.CreateColumnMajorLoad(
       ElemTy, $data, align, $stride, $isVolatile, $rows,
       $columns);
   }];
-  let assemblyFormat = "$data `,` `<` `stride` `=` $stride `>` attr-dict"
-    "`:` type($res) `from` type($data) `stride` type($stride)";
+  string mlirBuilder = [{
+    $res = $_builder.create<LLVM::MatrixColumnMajorLoadOp>(
+      $_location, $_resultType, $data, $stride,
+      $_int_attr($isVolatile), $_int_attr($rows), $_int_attr($columns));
+  }];
 }
 
 /// Create a column major, strided 2-D matrix store, as specified in the LLVM
 /// MatrixBuilder.
 /// matrix     - Matrix to store
 /// ptr        - Pointer to write back to
 /// isVolatile - True if the load operation is marked as volatile.
 /// rows       - Number of rows in matrix (must be a constant)
 /// columns    - Number of columns in matrix (must be a constant)
 /// stride     - Space between columns
-def LLVM_MatrixColumnMajorStoreOp : LLVM_Op<"intr.matrix.column.major.store"> {
+def LLVM_MatrixColumnMajorStoreOp : LLVM_ZeroResultIntrOp<"matrix.column.major.store"> {
   let arguments = (ins LLVM_AnyVector:$matrix, LLVM_Type:$data,
                    LLVM_Type:$stride, I1Attr:$isVolatile, I32Attr:$rows,
                    I32Attr:$columns);
   let builders = [LLVM_VoidResultTypeOpBuilder, LLVM_ZeroResultOpBuilder];
+  let assemblyFormat = "$matrix `,` $data `,` `<` `stride` `=` $stride `>` "
+    "attr-dict`:` type($matrix) `to` type($data) `stride` type($stride)";
+
   string llvmBuilder = [{
     llvm::MatrixBuilder mb(builder);
     const llvm::DataLayout &dl =
       builder.GetInsertBlock()->getModule()->getDataLayout();
     Type elementType = getVectorElementType(op.getMatrix().getType());
     llvm::Align align = dl.getABITypeAlign(
       moduleTranslation.convertType(elementType));
     mb.CreateColumnMajorStore(
       $matrix, $data, align, $stride, $isVolatile,
       $rows, $columns);
   }];
-  let assemblyFormat = "$matrix `,` $data `,` `<` `stride` `=` $stride `>` "
-    "attr-dict`:` type($matrix) `to` type($data) `stride` type($stride)";
+  string mlirBuilder = [{
+    $_builder.create<LLVM::MatrixColumnMajorStoreOp>(
+      $_location, $matrix, $data, $stride,
+      $_int_attr($isVolatile), $_int_attr($rows), $_int_attr($columns));
+  }];
 }
 
 /// Create a llvm.matrix.multiply call, multiplying 2-D matrices LHS and RHS, as
 /// specified in the LLVM MatrixBuilder.
-def LLVM_MatrixMultiplyOp : LLVM_Op<"intr.matrix.multiply"> {
+def LLVM_MatrixMultiplyOp : LLVM_OneResultIntrOp<"matrix.multiply"> {
   let arguments = (ins LLVM_Type:$lhs, LLVM_Type:$rhs, I32Attr:$lhs_rows,
                    I32Attr:$lhs_columns, I32Attr:$rhs_columns);
   let results = (outs LLVM_Type:$res);
   let builders = [LLVM_OneResultOpBuilder];
+  let assemblyFormat = "$lhs `,` $rhs attr-dict "
+    "`:` `(` type($lhs) `,` type($rhs) `)` `->` type($res)";
+
   string llvmBuilder = [{
     llvm::MatrixBuilder mb(builder);
     $res = mb.CreateMatrixMultiply(
       $lhs, $rhs, $lhs_rows, $lhs_columns,
       $rhs_columns);
   }];
-  let assemblyFormat = "$lhs `,` $rhs attr-dict "
-    "`:` `(` type($lhs) `,` type($rhs) `)` `->` type($res)";
+  string mlirBuilder = [{
+    $res = $_builder.create<LLVM::MatrixMultiplyOp>(
+      $_location, $_resultType, $lhs, $rhs,
+      $_int_attr($lhs_rows), $_int_attr($lhs_columns), $_int_attr($rhs_columns));
+  }];
 }
 
 /// Create a llvm.matrix.transpose call, transposing a `rows` x `columns` 2-D
 /// `matrix`, as specified in the LLVM MatrixBuilder.
-def LLVM_MatrixTransposeOp : LLVM_Op<"intr.matrix.transpose"> {
+def LLVM_MatrixTransposeOp : LLVM_OneResultIntrOp<"matrix.transpose"> {
   let arguments = (ins LLVM_Type:$matrix, I32Attr:$rows, I32Attr:$columns);
   let results = (outs LLVM_Type:$res);
   let builders = [LLVM_OneResultOpBuilder];
+  let assemblyFormat = "$matrix attr-dict `:` type($matrix) `into` type($res)";
+
   string llvmBuilder = [{
     llvm::MatrixBuilder mb(builder);
     $res = mb.CreateMatrixTranspose(
       $matrix, $rows, $columns);
   }];
-  let assemblyFormat = "$matrix attr-dict `:` type($matrix) `into` type($res)";
+  string mlirBuilder = [{
+    $res = $_builder.create<LLVM::MatrixTransposeOp>(
+      $_location, $_resultType, $matrix,
+      $_int_attr($rows), $_int_attr($columns));
+  }];
 }
 
 //
 // LLVM masked operations.
 //
 
 /// Create a llvm.get.active.lane.mask to set a mask up to a given position.
 def LLVM_GetActiveLaneMaskOp

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

 class LLVM_VectorReduction<string mnem>
     : LLVM_OneResultIntrOp<"vector.reduce." # mnem,
                            [], [0], [Pure]>,
       Arguments<(ins LLVM_Type)>;
 
 // LLVM vector reduction over a single vector, with an initial value,
 // and with permission to reassociate the reduction operations.
 class LLVM_VectorReductionAcc<string mnem>
-    : LLVM_OpBase<LLVM_Dialect, "intr.vector.reduce." # mnem,
-                  [Pure]>,
-      Results<(outs LLVM_Type:$res)>,
-      Arguments<(ins LLVM_Type, LLVM_Type,
+    : LLVM_OneResultIntrOp<"vector.reduce." # mnem,
+                           [], [0], [Pure]>,
+      Arguments<(ins LLVM_Type:$start_value, LLVM_Type:$input,
                  DefaultValuedAttr<BoolAttr, "false">:$reassoc)> {
   let llvmBuilder = [{
     llvm::Module *module = builder.GetInsertBlock()->getModule();
     llvm::Function *fn = llvm::Intrinsic::getDeclaration(
         module,
         llvm::Intrinsic::vector_reduce_}] # mnem # [{,
         { }] # !interleave(ListIntSubst<LLVM_IntrPatterns.operand, [1]>.lst,
                            ", ") # [{
         });
     auto operands = moduleTranslation.lookupValues(opInst.getOperands());
     llvm::FastMathFlags origFM = builder.getFastMathFlags();
     llvm::FastMathFlags tempFM = origFM;
     tempFM.setAllowReassoc($reassoc);
     builder.setFastMathFlags(tempFM);  // set fastmath flag
     $res = builder.CreateCall(fn, operands);
     builder.setFastMathFlags(origFM);  // restore fastmath flag
   }];
+  let mlirBuilder = [{
+    bool allowReassoc = inst->getFastMathFlags().allowReassoc();
+    $res = $_builder.create<$_qualCppClassName>($_location,
+      $_resultType, $start_value, $input, allowReassoc);
+  }];
 }
 
 def LLVM_OneResultOpBuilder :
   OpBuilder<(ins "Type":$resultType, "ValueRange":$operands,
     CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
   [{
     if (resultType) $_state.addTypes(resultType);
     $_state.addOperands(operands);

mlir/test/Target/LLVMIR/Import/intrinsic.ll

 ; CHECK-LABEL:  llvm.func @umin_test
 define void @umin_test(i32 %0, i32 %1, <8 x i32> %2, <8 x i32> %3) {
   ; CHECK:   "llvm.intr.umin"(%{{.*}}, %{{.*}}) : (i32, i32) -> i32
   %5 = call i32 @llvm.umin.i32(i32 %0, i32 %1)
   ; CHECK:   "llvm.intr.umin"(%{{.*}}, %{{.*}}) : (vector<8xi32>, vector<8xi32>) -> vector<8xi32>
   %6 = call <8 x i32> @llvm.umin.v8i32(<8 x i32> %2, <8 x i32> %3)
   ret void
 }
+
 ; CHECK-LABEL:  llvm.func @vector_reductions
 define void @vector_reductions(float %0, <8 x float> %1, <8 x i32> %2) {
   ; CHECK: "llvm.intr.vector.reduce.add"(%{{.*}}) : (vector<8xi32>) -> i32
   %4 = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> %2)
   ; CHECK: "llvm.intr.vector.reduce.and"(%{{.*}}) : (vector<8xi32>) -> i32
   %5 = call i32 @llvm.vector.reduce.and.v8i32(<8 x i32> %2)
   ; CHECK: "llvm.intr.vector.reduce.fmax"(%{{.*}}) : (vector<8xf32>) -> f32
   %6 = call float @llvm.vector.reduce.fmax.v8f32(<8 x float> %1)
   ; CHECK: "llvm.intr.vector.reduce.fmin"(%{{.*}}) : (vector<8xf32>) -> f32
   %7 = call float @llvm.vector.reduce.fmin.v8f32(<8 x float> %1)
   ; CHECK: "llvm.intr.vector.reduce.mul"(%{{.*}}) : (vector<8xi32>) -> i32
   %8 = call i32 @llvm.vector.reduce.mul.v8i32(<8 x i32> %2)
   ; CHECK: "llvm.intr.vector.reduce.or"(%{{.*}}) : (vector<8xi32>) -> i32
   %9 = call i32 @llvm.vector.reduce.or.v8i32(<8 x i32> %2)
   ; CHECK: "llvm.intr.vector.reduce.smax"(%{{.*}}) : (vector<8xi32>) -> i32
   %10 = call i32 @llvm.vector.reduce.smax.v8i32(<8 x i32> %2)
   ; CHECK: "llvm.intr.vector.reduce.smin"(%{{.*}}) : (vector<8xi32>) -> i32
   %11 = call i32 @llvm.vector.reduce.smin.v8i32(<8 x i32> %2)
   ; CHECK: "llvm.intr.vector.reduce.umax"(%{{.*}}) : (vector<8xi32>) -> i32
   %12 = call i32 @llvm.vector.reduce.umax.v8i32(<8 x i32> %2)
   ; CHECK: "llvm.intr.vector.reduce.umin"(%{{.*}}) : (vector<8xi32>) -> i32
   %13 = call i32 @llvm.vector.reduce.umin.v8i32(<8 x i32> %2)
-  ; TODO: vector reduce fadd and fmul should be handled specially.
+  ; CHECK: "llvm.intr.vector.reduce.fadd"(%{{.*}}, %{{.*}}) {reassoc = false} : (f32, vector<8xf32>) -> f32
   %14 = call float @llvm.vector.reduce.fadd.v8f32(float %0, <8 x float> %1)
+  ; CHECK: "llvm.intr.vector.reduce.fmul"(%{{.*}}, %{{.*}}) {reassoc = false} : (f32, vector<8xf32>) -> f32
   %15 = call float @llvm.vector.reduce.fmul.v8f32(float %0, <8 x float> %1)
+  ; CHECK: "llvm.intr.vector.reduce.fadd"(%{{.*}}, %{{.*}}) {reassoc = true} : (f32, vector<8xf32>) -> f32
   %16 = call reassoc float @llvm.vector.reduce.fadd.v8f32(float %0, <8 x float> %1)
+  ; CHECK: "llvm.intr.vector.reduce.fmul"(%{{.*}}, %{{.*}}) {reassoc = true} : (f32, vector<8xf32>) -> f32
   %17 = call reassoc float @llvm.vector.reduce.fmul.v8f32(float %0, <8 x float> %1)
   ; CHECK:  "llvm.intr.vector.reduce.xor"(%{{.*}}) : (vector<8xi32>) -> i32
   %18 = call i32 @llvm.vector.reduce.xor.v8i32(<8 x i32> %2)
   ret void
 }
 
-; TODO: matrix intrinsic should be handled specially.
-define void @matrix_intrinsics(<64 x float> %0, <48 x float> %1, float* %2, i64 %3) {
-  %5 = call <12 x float> @llvm.matrix.multiply.v12f32.v64f32.v48f32(<64 x float> %0, <48 x float> %1, i32 4, i32 16, i32 3)
-  %6 = call <48 x float> @llvm.matrix.transpose.v48f32(<48 x float> %1, i32 3, i32 16)
-  %7 = call <48 x float> @llvm.matrix.column.major.load.v48f32.i64(float* align 4 %2, i64 %3, i1 false, i32 3, i32 16)
-  call void @llvm.matrix.column.major.store.v48f32.i64(<48 x float> %7, float* align 4 %2, i64 %3, i1 false, i32 3, i32 16)
+; CHECK-LABEL: @matrix_intrinsics
+; CHECK-SAME:  %[[VEC1:[a-zA-Z0-9]+]]
+; CHECK-SAME:  %[[VEC2:[a-zA-Z0-9]+]]
+; CHECK-SAME:  %[[PTR:[a-zA-Z0-9]+]]
+; CHECK-SAME:  %[[STRIDE:[a-zA-Z0-9]+]]
+define void @matrix_intrinsics(<64 x float> %vec1, <48 x float> %vec2, float* %ptr, i64 %stride) {
+  ; CHECK:  llvm.intr.matrix.multiply %[[VEC1]], %[[VEC2]]
+  ; CHECK-SAME:  {lhs_columns = 16 : i32, lhs_rows = 4 : i32, rhs_columns = 3 : i32}
+  %1 = call <12 x float> @llvm.matrix.multiply.v12f32.v64f32.v48f32(<64 x float> %vec1, <48 x float> %vec2, i32 4, i32 16, i32 3)
+  ; CHECK:  llvm.intr.matrix.transpose %[[VEC2]]
+  ; CHECK-SAME:  {columns = 16 : i32, rows = 3 : i32}
+  %2 = call <48 x float> @llvm.matrix.transpose.v48f32(<48 x float> %vec2, i32 3, i32 16)
+  ; CHECK:  %[[VAL1:.+]] = llvm.intr.matrix.column.major.load %[[PTR]], <stride = %[[STRIDE]]>
+  ; CHECK-SAME:  {columns = 16 : i32, isVolatile = false, rows = 3 : i32}
+  %3 = call <48 x float> @llvm.matrix.column.major.load.v48f32.i64(float* align 4 %ptr, i64 %stride, i1 false, i32 3, i32 16)
+  ; CHECK:  llvm.intr.matrix.column.major.store %[[VAL1]], %[[PTR]], <stride = %[[STRIDE]]>
+  ; CHECK-SAME:  {columns = 16 : i32, isVolatile = true, rows = 3 : i32}
+  call void @llvm.matrix.column.major.store.v48f32.i64(<48 x float> %3, float* align 4 %ptr, i64 %stride, i1 true, i32 3, i32 16)
   ret void
 }
 
 ; CHECK-LABEL: llvm.func @get_active_lane_mask
 define <7 x i1> @get_active_lane_mask(i64 %0, i64 %1) {
   ; CHECK: llvm.intr.get.active.lane.mask %{{.*}}, %{{.*}} : i64, i64 to vector<7xi1>
   %3 = call <7 x i1> @llvm.get.active.lane.mask.v7i1.i64(i64 %0, i64 %1)
   ret <7 x i1> %3
; CHECK: llvm.intr.vastart %{{.*}}
   call void @llvm.va_start(i8* %0)
 ; CHECK: llvm.intr.vacopy %{{.*}} to %{{.*}}
   call void @llvm.va_copy(i8* %1, i8* %0)
 ; CHECK: llvm.intr.vaend %{{.*}}
   call void @llvm.va_end(i8* %0)
   ret void
 }
 
+; CHECK-LABEL: @assume
+; CHECK-SAME:  %[[TRUE:[a-zA-Z0-9]+]]
+define void @assume(i1 %true) {
+  ; CHECK:  "llvm.intr.assume"(%[[TRUE]]) : (i1) -> ()
+  call void @llvm.assume(i1 %true)
+  ret void
+}
+
 ; CHECK-LABEL:  llvm.func @coro_id
 define void @coro_id(i32 %0, i8* %1) {
   ; CHECK: llvm.intr.coro.id %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : !llvm.token
   %3 = call token @llvm.coro.id(i32 %0, i8* %1, i8* %1, i8* null)
   ret void
 }
 
 ; CHECK-LABEL:  llvm.func @coro_begin

 declare <8 x float> @llvm.vp.fptrunc.v8f32.v8f64(<8 x double>, <8 x i1>, i32)
 declare <8 x double> @llvm.vp.fpext.v8f64.v8f32(<8 x float>, <8 x i1>, i32)
 declare <8 x i64> @llvm.vp.fptoui.v8i64.v8f64(<8 x double>, <8 x i1>, i32)
 declare <8 x i64> @llvm.vp.fptosi.v8i64.v8f64(<8 x double>, <8 x i1>, i32)
 declare <8 x i64> @llvm.vp.ptrtoint.v8i64.v8p0i32(<8 x i32*>, <8 x i1>, i32)
 declare <8 x i32*> @llvm.vp.inttoptr.v8p0i32.v8i64(<8 x i64>, <8 x i1>, i32)
 declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture)
 declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture)
+declare void @llvm.assume(i1)