[MLIR] Avoid LLVM pointer element type accesses in LLVMOps.td

[nikic]

Summary

For compatibility with opaque pointers, we should be using the MLIR-level pointer element type (and translating it to LLVM), not the LLVM-level pointer element type.

[ftynse]

I don't think this is the right approach. If LLVM fully adopts opaque pointers, so should the LLVM dialect in MLIR. Relying on the LLVM dialect having access to element types of a pointer through MLIR types only displaces the problem.

[nikic]

In D122708#3416697, @ftynse wrote:

I don't think this is the right approach. If LLVM fully adopts opaque pointers, so should the LLVM dialect in MLIR. Relying on the LLVM dialect having access to element types of a pointer through MLIR types only displaces the problem.

Oh, that's interesting. I assumed it was a deliberate design decision to use typed pointer on the MLIR level, as there's no strong requirement for these to match, and typed pointers might be easier to work with on a higher level of abstraction. Of course, if MLIR would like to switch to opaque pointers as well, that sounds great to me.

I'm not familiar with MLIR myself, so I don't know what moving MLIR to opaque pointers would involve in practical terms.

[mehdi_amini]

There is a (very incomplete) import path from LLVM IR to the LLVM dialect, I assume this will make it a requirement for the LLVM Dialect to adopt typeless pointer "at some point" anyway right?
Even if we had the choice, the role of the LLVM dialect in MLIR is to stay as a close as possible to LLVM IR to keep the translation between the two "trivial".

[nikic]

In D122708#3419852, @mehdi_amini wrote:

There is a (very incomplete) import path from LLVM IR to the LLVM dialect, I assume this will make it a requirement for the LLVM Dialect to adopt typeless pointer "at some point" anyway right?

Not necessarily, it's possible to create use a dummy pointer type and bitcast it appropriately for certain uses. Wouldn't be very pretty, of course.

Even if we had the choice, the role of the LLVM dialect in MLIR is to stay as a close as possible to LLVM IR to keep the translation between the two "trivial".

That makes sense.

[ftynse]

I will look into adding support for opaque types in the LLVM dialect, and then this can be rebased.

[nikic]

MLIR support for opaque pointers was added in https://reviews.llvm.org/D123310 and https://reviews.llvm.org/D123792, so abandoning this revision.

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

 }
 
 // Memory-related operations.
 def LLVM_AllocaOp : LLVM_Op<"alloca">, MemoryOpWithAlignmentBase {
   let arguments = (ins AnyInteger:$arraySize,
                    OptionalAttr<I64Attr>:$alignment);
   let results = (outs Res<LLVM_AnyPointer, "", [MemAlloc<AutomaticAllocationScopeResource>]>:$res);
   string llvmBuilder = [{
+    auto ty = op.getRes().getType().cast<LLVMPointerType>().getElementType();
     auto addrSpace = $_resultType->getPointerAddressSpace();
     auto *inst = builder.CreateAlloca(
-        $_resultType->getPointerElementType(), addrSpace, $arraySize);
+        moduleTranslation.convertType(ty), addrSpace, $arraySize);
     }] # setAlignmentCode # [{
     $res = inst;
   }];
   let builders = [
     OpBuilder<(ins "Type":$resultType, "Value":$arraySize,
       "unsigned":$alignment),
     [{
       if (alignment == 0)
let llvmBuilder = [{
     indices.reserve($structIndices.size());
     unsigned operandIdx = 0;
     for (int32_t structIndex : $structIndices.getValues<int32_t>()) {
       if (structIndex == GEPOp::kDynamicIndex)
         indices.push_back($indices[operandIdx++]);
       else
         indices.push_back(builder.getInt32(structIndex));
     }
+    auto ptrTy = op.getBase().getType();
+    if (LLVM::isCompatibleVectorType(ptrTy))
+      ptrTy = LLVM::getVectorElementType(ptrTy);
+    auto elemTy = ptrTy.cast<LLVMPointerType>().getElementType();
     $res = builder.CreateGEP(
-      $base->getType()->getPointerElementType(), $base, indices);
+        moduleTranslation.convertType(elemTy), $base, indices);
   }];
   let assemblyFormat = [{
     $base `[` custom<GEPIndices>($indices, $structIndices) `]` attr-dict
     `:` functional-type(operands, results)
   }];
 
   let extraClassDeclaration = [{
     constexpr static int kDynamicIndex = std::numeric_limits<int32_t>::min();
   }];
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 def LLVM_LoadOp : LLVM_Op<"load">, MemoryOpWithAlignmentAndAttributes {
   let arguments = (ins Arg<LLVM_PointerTo<LLVM_LoadableType>, "", [MemRead]>:$addr,
                    OptionalAttr<SymbolRefArrayAttr>:$access_groups,
                    OptionalAttr<SymbolRefArrayAttr>:$alias_scopes,
                    OptionalAttr<SymbolRefArrayAttr>:$noalias_scopes,
                    OptionalAttr<I64Attr>:$alignment, UnitAttr:$volatile_,
                    UnitAttr:$nontemporal);
   let results = (outs LLVM_LoadableType:$res);
   string llvmBuilder = [{
+    auto accessTy =
+        op.getAddr().getType().cast<LLVMPointerType>().getElementType();
     auto *inst = builder.CreateLoad(
-      $addr->getType()->getPointerElementType(), $addr, $volatile_);
+        moduleTranslation.convertType(accessTy), $addr, $volatile_);
   }] # setAlignmentCode
      # setNonTemporalMetadataCode
      # setAccessGroupsMetadataCode
      # setAliasScopeMetadataCode
      # [{
     $res = inst;
   }];
   let builders = [

 /// 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"> {
-  let arguments = (ins LLVM_Type:$data, LLVM_Type:$stride, I1Attr:$isVolatile,
-                   I32Attr:$rows, I32Attr:$columns);
+  let arguments = (ins LLVM_AnyPointer:$data, LLVM_Type:$stride,
+                   I1Attr:$isVolatile, I32Attr:$rows, I32Attr:$columns);
   let results = (outs LLVM_Type:$res);
   let builders = [LLVM_OneResultOpBuilder];
   string llvmBuilder = [{
     llvm::MatrixBuilder mb(builder);
     const llvm::DataLayout &dl =
       builder.GetInsertBlock()->getModule()->getDataLayout();
-    llvm::Type *ElemTy = $data->getType()->getPointerElementType();
-    llvm::Align align = dl.getABITypeAlign(ElemTy);
+    auto ty = op.getData().getType().cast<LLVMPointerType>().getElementType();
+    llvm::Type *llvmTy = moduleTranslation.convertType(ty);
+    llvm::Align align = dl.getABITypeAlign(llvmTy);
     $res = mb.CreateColumnMajorLoad(
-      ElemTy, $data, align, $stride, $isVolatile, $rows,
-      $columns);
+      llvmTy, $data, align, $stride, $isVolatile, $rows, $columns);
   }];
   let assemblyFormat = "$data `,` `<` `stride` `=` $stride `>` attr-dict"
     "`:` type($res) `from` type($data) `stride` type($stride)";
 }
 
 /// 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"> {
-  let arguments = (ins LLVM_Type:$matrix, LLVM_Type:$data, LLVM_Type:$stride,
-                   I1Attr:$isVolatile, I32Attr:$rows, I32Attr:$columns);
+  let arguments = (ins LLVM_Type:$matrix, LLVM_AnyPointer:$data,
+                   LLVM_Type:$stride, I1Attr:$isVolatile, I32Attr:$rows,
+                   I32Attr:$columns);
   let builders = [LLVM_VoidResultTypeOpBuilder, LLVM_ZeroResultOpBuilder];
   string llvmBuilder = [{
     llvm::MatrixBuilder mb(builder);
     const llvm::DataLayout &dl =
       builder.GetInsertBlock()->getModule()->getDataLayout();
-    llvm::Align align = dl.getABITypeAlign(
-      $data->getType()->getPointerElementType());
+    auto ty = op.getData().getType().cast<LLVMPointerType>().getElementType();
+    llvm::Type *llvmTy = moduleTranslation.convertType(ty);
+    llvm::Align align = dl.getABITypeAlign(llvmTy);
     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)";
 }
 
def LLVM_GetActiveLaneMaskOp
     : LLVM_OneResultIntrOp<"get.active.lane.mask", [0], [0], [NoSideEffect]> {
   let arguments = (ins LLVM_Type:$base, LLVM_Type:$n);
   let assemblyFormat = "$base `,` $n attr-dict `:` "
     "type($base) `,` type($n) `to` type($res)";
 }
 
 /// Create a call to Masked Load intrinsic.
 def LLVM_MaskedLoadOp : LLVM_Op<"intr.masked.load"> {
-  let arguments = (ins LLVM_Type:$data, LLVM_Type:$mask,
+  let arguments = (ins LLVM_AnyPointer:$data, LLVM_Type:$mask,
                    Variadic<LLVM_Type>:$pass_thru, I32Attr:$alignment);
   let results = (outs LLVM_Type:$res);
   let builders = [LLVM_OneResultOpBuilder];
   string llvmBuilder = [{
-    llvm::Type *Ty = $data->getType()->getPointerElementType();
+    auto ty = op.getData().getType().cast<LLVMPointerType>().getElementType();
+    llvm::Type *llvmTy = moduleTranslation.convertType(ty);
     $res = $pass_thru.empty() ? builder.CreateMaskedLoad(
-      Ty, $data, llvm::Align($alignment), $mask) :
+      llvmTy, $data, llvm::Align($alignment), $mask) :
       builder.CreateMaskedLoad(
-        Ty, $data, llvm::Align($alignment), $mask, $pass_thru[0]);
+        llvmTy, $data, llvm::Align($alignment), $mask, $pass_thru[0]);
   }];
   let assemblyFormat =
     "operands attr-dict `:` functional-type(operands, results)";
 }
 
 /// Create a call to Masked Store intrinsic.
 def LLVM_MaskedStoreOp : LLVM_Op<"intr.masked.store"> {
   let arguments = (ins LLVM_Type:$value, LLVM_Type:$data, LLVM_Type:$mask,
                    I32Attr:$alignment);
   let builders = [LLVM_VoidResultTypeOpBuilder, LLVM_ZeroResultOpBuilder];
   string llvmBuilder = [{
     builder.CreateMaskedStore(
       $value, $data, llvm::Align($alignment), $mask);
   }];
   let assemblyFormat = "$value `,` $data `,` $mask attr-dict `:` "
     "type($value) `,` type($mask) `into` type($data)";
 }
 
 /// Create a call to Masked Gather intrinsic.
 def LLVM_masked_gather : LLVM_Op<"intr.masked.gather"> {
-  let arguments = (ins LLVM_Type:$ptrs, LLVM_Type:$mask,
+  let arguments = (ins LLVM_VectorOf<LLVM_AnyPointer>:$ptrs, LLVM_Type:$mask,
                    Variadic<LLVM_Type>:$pass_thru, I32Attr:$alignment);
   let results = (outs LLVM_Type:$res);
   let builders = [LLVM_OneResultOpBuilder];
   string llvmBuilder = [{
-    llvm::VectorType *PtrVecTy = cast<llvm::VectorType>($ptrs->getType());
+    auto vecTy = op.getPtrs().getType();
+    auto ptrTy = LLVM::getVectorElementType(vecTy).cast<LLVMPointerType>();
     llvm::Type *Ty = llvm::VectorType::get(
-      PtrVecTy->getElementType()->getPointerElementType(),
-      PtrVecTy->getElementCount());
+      moduleTranslation.convertType(ptrTy.getElementType()),
+      LLVM::getVectorNumElements(vecTy));
     $res = $pass_thru.empty() ? builder.CreateMaskedGather(
       Ty, $ptrs, llvm::Align($alignment), $mask) :
       builder.CreateMaskedGather(
         Ty, $ptrs, llvm::Align($alignment), $mask, $pass_thru[0]);
   }];
   let assemblyFormat =
     "operands attr-dict `:` functional-type(operands, results)";
 }