[mlir] [LLVMIR] Add get active lane mask intrinsic

[aartbik]

Summary

Provides fast, generic way of setting a mask up to a certain
point. Potential use cases that may benefit are create_mask
and transfer_read/write operations in the vector dialect.

[bkramer]

Can we use LLVM_OneResultIntrOp to simplify the .td boilerplate? Looks good otherwise.

[aartbik]

addressed comments

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

def LLVM_MatrixTransposeOp
   }];
   let assemblyFormat = "$matrix attr-dict `:` type($matrix) `into` type($res)";
 }
 
 //
 // LLVM masked operations.
 //
 
+/// Create a llvm.get.active.lane.mask to set a mask up to a given position.
+def LLVM_GetActiveLaneMaskOp
+    : LLVM_OneResultOp<"intr.get.active.lane.mask">,
+      Arguments<(ins LLVM_Type:$base, LLVM_Type:$n)> {
+  let llvmBuilder = [{
+    llvm::Module *module = builder.GetInsertBlock()->getModule();
+    llvm::Function *fn = llvm::Intrinsic::getDeclaration(
+        module,
+        llvm::Intrinsic::get_active_lane_mask,
+        { }] # StrJoin<!listconcat(
+            ListIntSubst<LLVM_IntrPatterns.result, [0]>.lst,
+            ListIntSubst<LLVM_IntrPatterns.operand, [0]>.lst)>.result # [{
+        });
+    auto operands = lookupValues(opInst.getOperands());
+    $res = builder.CreateCall(fn, operands);
+  }];
+  let assemblyFormat = "$base `,` $n attr-dict `:` "
+    "type($base) `,` type($n) `to` type($res)";
+}
+
 /// Create a call to Masked Load intrinsic.
 def LLVM_MaskedLoadOp
     : LLVM_OneResultOp<"intr.masked.load">,
       Arguments<(ins LLVM_Type:$data, LLVM_Type:$mask,
                  Variadic<LLVM_Type>:$pass_thru, I32Attr:$alignment)> {
   string llvmBuilder = [{
     $res = $pass_thru.empty() ? builder.CreateMaskedLoad(
       $data, llvm::Align($alignment.getZExtValue()), $mask) :

mlir/test/Target/llvmir-intrinsics.mlir

%E = llvm.intr.matrix.column.major.load %ptr, <stride=%stride>
     !llvm.vec<48 x float> from !llvm.ptr<float> stride !llvm.i64
   // CHECK: call void @llvm.matrix.column.major.store.v48f32(<48 x float> %7, float* align 4 %2, i64 %3, i1 false, i32 3, i32 16)
   llvm.intr.matrix.column.major.store %E, %ptr, <stride=%stride>
     { isVolatile = 0: i1, rows = 3: i32, columns = 16: i32} :
     !llvm.vec<48 x float> to !llvm.ptr<float> stride !llvm.i64
   llvm.return
 }
 
+// CHECK-LABEL: @get_active_lane_mask
+llvm.func @get_active_lane_mask(%base: !llvm.i64, %n: !llvm.i64) -> (!llvm.vec<7 x i1>) {
+  // CHECK: call <7 x i1> @llvm.get.active.lane.mask.v7i1.i64(i64 %0, i64 %1)
+  %0 = llvm.intr.get.active.lane.mask %base, %n : !llvm.i64, !llvm.i64 to !llvm.vec<7 x i1>
+  llvm.return %0 : !llvm.vec<7 x i1>
+}
+
 // CHECK-LABEL: @masked_load_store_intrinsics
 llvm.func @masked_load_store_intrinsics(%A: !llvm.ptr<vec<7 x float>>, %mask: !llvm.vec<7 x i1>) {
   // CHECK: call <7 x float> @llvm.masked.load.v7f32.p0v7f32(<7 x float>* %{{.*}}, i32 1, <7 x i1> %{{.*}}, <7 x float> undef)
   %a = llvm.intr.masked.load %A, %mask { alignment = 1: i32} :
     (!llvm.ptr<vec<7 x float>>, !llvm.vec<7 x i1>) -> !llvm.vec<7 x float>
   // CHECK: call <7 x float> @llvm.masked.load.v7f32.p0v7f32(<7 x float>* %{{.*}}, i32 1, <7 x i1> %{{.*}}, <7 x float> %{{.*}})
   %b = llvm.intr.masked.load %A, %mask, %a { alignment = 1: i32} :
     (!llvm.ptr<vec<7 x float>>, !llvm.vec<7 x i1>, !llvm.vec<7 x float>) -> !llvm.vec<7 x float>

 // CHECK-DAG: declare <8 x float> @llvm.ceil.v8f32(<8 x float>) #0
 // CHECK-DAG: declare float @llvm.cos.f32(float)
 // CHECK-DAG: declare <8 x float> @llvm.cos.v8f32(<8 x float>) #0
 // CHECK-DAG: declare float @llvm.copysign.f32(float, float)
 // CHECK-DAG: declare <12 x float> @llvm.matrix.multiply.v12f32.v64f32.v48f32(<64 x float>, <48 x float>, i32 immarg, i32 immarg, i32 immarg)
 // CHECK-DAG: declare <48 x float> @llvm.matrix.transpose.v48f32(<48 x float>, i32 immarg, i32 immarg)
 // CHECK-DAG: declare <48 x float> @llvm.matrix.column.major.load.v48f32(float* nocapture, i64, i1 immarg, i32 immarg, i32 immarg)
 // CHECK-DAG: declare void @llvm.matrix.column.major.store.v48f32(<48 x float>, float* nocapture writeonly, i64, i1 immarg, i32 immarg, i32 immarg)
+// CHECK-DAG: declare <7 x i1> @llvm.get.active.lane.mask.v7i1.i64(i64, i64)
 // CHECK-DAG: declare <7 x float> @llvm.masked.load.v7f32.p0v7f32(<7 x float>*, i32 immarg, <7 x i1>, <7 x float>)
 // CHECK-DAG: declare void @llvm.masked.store.v7f32.p0v7f32(<7 x float>, <7 x float>*, i32 immarg, <7 x i1>)
+// CHECK-DAG: declare <7 x float> @llvm.masked.gather.v7f32.v7p0f32(<7 x float*>, i32 immarg, <7 x i1>, <7 x float>)
+// CHECK-DAG: declare void @llvm.masked.scatter.v7f32.v7p0f32(<7 x float>, <7 x float*>, i32 immarg, <7 x i1>)
+// CHECK-DAG: declare <7 x float> @llvm.masked.expandload.v7f32(float*, <7 x i1>, <7 x float>)
+// CHECK-DAG: declare void @llvm.masked.compressstore.v7f32(<7 x float>, float*, <7 x i1>)
 // CHECK-DAG: declare void @llvm.memcpy.p0i8.p0i8.i32(i8* noalias nocapture writeonly, i8* noalias nocapture readonly, i32, i1 immarg)
 // CHECK-DAG: declare void @llvm.memcpy.inline.p0i8.p0i8.i64(i8* noalias nocapture writeonly, i8* noalias nocapture readonly, i64 immarg, i1 immarg)