[mlir][llvm] Add inbounds attriubte to the gep op.

[gysit]

Summary

The revision adds an inbounds attribute to the LLVM dialect
GEP operation. It extends the builders and the import and export
to support the optional inbounds attribute.

As all builders set inbounds to false by default, existing lowerings
from higher-level dialects to LLVM dialect are not affected by the
change. Canonicalization/folding remains untouched since it currently
does not implement any simplifications in case of undefined behavior
(the handling of undefined behavior is deferred to LLVM).

[ftynse]

Thanks for doing what should have been done years ago! :)

I've been thinking about "magic bools" in builders (b.create<GEPOp>(..., /*inbounds=*/true)). What's your opinion on using named zero-size structs instead (b.create<GEPOp>(..., GEPOp::InBounds()))? Not asking to implement, just for an opinion.

[gysit]

I've been thinking about "magic bools" in builders (b.create<GEPOp>(..., /*inbounds=*/true)). What's your opinion on using named zero-size structs instead (b.create<GEPOp>(..., GEPOp::InBounds()))? Not asking to implement, just for an opinion.

I like the idea. Having an explicit GEPOp::InBounds() is clearly more readable and trustworthy than the c-style comments /*inbounds=*/true. Especially, since I tend to look at other builder uses due to the missing autocompletion (I know I should look at the tablegen file :)).

How would you handle default values with this design? Would you provide multiple build methods instead of having a default value or is there another trick?

[ftynse]

Multiple overloads, or wrapping a bool into a struct with the default constructor putting a true and a regular taking the value that is set to false in the default argument initialization.

[gysit]

Multiple overloads, or wrapping a bool into a struct with the default constructor putting a true and a regular taking the value that is set to false in the default argument initialization.

Yeah, I think wrapping a bool may be nice!

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

def LLVM_AllocaOp : LLVM_Op<"alloca">, MemoryOpWithAlignmentBase {
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 def LLVM_GEPOp : LLVM_Op<"getelementptr", [Pure]> {
   let arguments = (ins LLVM_ScalarOrVectorOf<LLVM_AnyPointer>:$base,
                    Variadic<LLVM_ScalarOrVectorOf<AnyInteger>>:$dynamicIndices,
                    DenseI32ArrayAttr:$rawConstantIndices,
-                   OptionalAttr<TypeAttr>:$elem_type);
+                   OptionalAttr<TypeAttr>:$elem_type,
+                   UnitAttr:$inbounds);
   let results = (outs LLVM_ScalarOrVectorOf<LLVM_AnyPointer>:$res);
   let skipDefaultBuilders = 1;
 
   let description = [{
     This operation mirrors LLVM IRs 'getelementptr' operation that is used to
     perform pointer arithmetic.
 
     Like in LLVM IR, it is possible to use both constants as well as SSA values
     as indices. In the case of indexing within a structure, it is required to
     either use constant indices directly, or supply a constant SSA value.
 
+    An optional 'inbounds' attribute specifies the low-level pointer arithmetic
+    overflow behavior that LLVM uses after lowering the operation to LLVM IR.
+
     Examples:
 
     ```mlir
     // GEP with an SSA value offset
     %0 = llvm.getelementptr %1[%2] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
 
-    // GEP with a constant offset
-    %0 = llvm.getelementptr %1[3] : (!llvm.ptr<f32>) -> !llvm.ptr<f32>
+    // GEP with a constant offset and the inbounds attribute set
+    %0 = llvm.getelementptr inbounds %1[3] : (!llvm.ptr<f32>) -> !llvm.ptr<f32>
 
     // GEP with constant offsets into a structure
     %0 = llvm.getelementptr %1[0, 1]
        : (!llvm.ptr<struct(i32, f32)>) -> !llvm.ptr<f32>
     ```
   }];
 
   let builders = [
     OpBuilder<(ins "Type":$resultType, "Type":$basePtrType, "Value":$basePtr,
-               "ValueRange":$indices,
+               "ValueRange":$indices, CArg<"bool", "false">:$inbounds,
                CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
-    OpBuilder<(ins "Type":$resultType, "Value":$basePtr, "ValueRange":$indices,
+    OpBuilder<(ins "Type":$resultType, "Value":$basePtr,
+               "ValueRange":$indices, CArg<"bool", "false">:$inbounds,
                CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
-    OpBuilder<(ins "Type":$resultType, "Value":$basePtr, "ArrayRef<GEPArg>":$indices,
+    OpBuilder<(ins "Type":$resultType, "Value":$basePtr,
+               "ArrayRef<GEPArg>":$indices, CArg<"bool", "false">:$inbounds,
                CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
     OpBuilder<(ins "Type":$resultType, "Type":$basePtrType, "Value":$basePtr,
-               "ArrayRef<GEPArg>":$indices,
+               "ArrayRef<GEPArg>":$indices, CArg<"bool", "false">:$inbounds,
                CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
   ];
   let llvmBuilder = [{
     SmallVector<llvm::Value *> indices;
     indices.reserve($rawConstantIndices.size());
     GEPIndicesAdaptor<decltype($dynamicIndices)>
         gepIndices(op.getRawConstantIndicesAttr(), $dynamicIndices);
     for (PointerUnion<IntegerAttr, llvm::Value*> valueOrAttr : gepIndices) {
       if (llvm::Value* value = valueOrAttr.dyn_cast<llvm::Value*>())
         indices.push_back(value);
       else
         indices.push_back(
             builder.getInt32(valueOrAttr.get<IntegerAttr>().getInt()));
     }
     Type baseElementType = op.getSourceElementType();
     llvm::Type *elementType = moduleTranslation.convertType(baseElementType);
-    $res = builder.CreateGEP(elementType, $base, indices);
+    $res = builder.CreateGEP(elementType, $base, indices, "", $inbounds);
   }];
   let assemblyFormat = [{
+    (`inbounds` $inbounds^)?
     $base `[` custom<GEPIndices>($dynamicIndices, $rawConstantIndices) `]` attr-dict
     `:` functional-type(operands, results) (`,` $elem_type^)?
   }];
 
   let extraClassDeclaration = [{
     constexpr static int32_t kDynamicIndex = std::numeric_limits<int32_t>::min();
 
     /// Returns the type pointed to by the pointer argument of this GEP.

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

static Type extractVectorElementType(Type type) {
   if (auto scalableVectorType = type.dyn_cast<LLVMScalableVectorType>())
     return scalableVectorType.getElementType();
   if (auto fixedVectorType = type.dyn_cast<LLVMFixedVectorType>())
     return fixedVectorType.getElementType();
   return type;
 }
 
 void GEPOp::build(OpBuilder &builder, OperationState &result, Type resultType,
-                  Value basePtr, ArrayRef<GEPArg> indices,
+                  Value basePtr, ArrayRef<GEPArg> indices, bool inbounds,
                   ArrayRef<NamedAttribute> attributes) {
   auto ptrType =
       extractVectorElementType(basePtr.getType()).cast<LLVMPointerType>();
   assert(!ptrType.isOpaque() &&
          "expected non-opaque pointer, provide elementType explicitly when "
          "opaque pointers are used");
   build(builder, result, resultType, ptrType.getElementType(), basePtr, indices,
-        attributes);
+        inbounds, attributes);
 }
 
 /// Destructures the 'indices' parameter into 'rawConstantIndices' and
 /// 'dynamicIndices', encoding the former in the process. In the process,
 /// dynamic indices which are used to index into a structure type are converted
 /// to constant indices when possible. To do this, the GEPs element type should
 /// be passed as first parameter.
 static void destructureIndices(Type currType, ArrayRef<GEPArg> indices,
currType =
               return nullptr;
             })
             .Default(Type(nullptr));
   }
 }
 
 void GEPOp::build(OpBuilder &builder, OperationState &result, Type resultType,
                   Type elementType, Value basePtr, ArrayRef<GEPArg> indices,
-                  ArrayRef<NamedAttribute> attributes) {
+                  bool inbounds, ArrayRef<NamedAttribute> attributes) {
   SmallVector<int32_t> rawConstantIndices;
   SmallVector<Value> dynamicIndices;
   destructureIndices(elementType, indices, rawConstantIndices, dynamicIndices);
 
   result.addTypes(resultType);
   result.addAttributes(attributes);
   result.addAttribute(getRawConstantIndicesAttrName(result.name),
                       builder.getDenseI32ArrayAttr(rawConstantIndices));
+  if (inbounds) {
+    result.addAttribute(getInboundsAttrName(result.name),
+                        builder.getUnitAttr());
+  }
   if (extractVectorElementType(basePtr.getType())
           .cast<LLVMPointerType>()
           .isOpaque())
     result.addAttribute(kElemTypeAttrName, TypeAttr::get(elementType));
   result.addOperands(basePtr);
   result.addOperands(dynamicIndices);
 }
 
 void GEPOp::build(OpBuilder &builder, OperationState &result, Type resultType,
-                  Value basePtr, ValueRange indices,
+                  Value basePtr, ValueRange indices, bool inbounds,
                   ArrayRef<NamedAttribute> attributes) {
   build(builder, result, resultType, basePtr, SmallVector<GEPArg>(indices),
-        attributes);
+        inbounds, attributes);
 }
 
 void GEPOp::build(OpBuilder &builder, OperationState &result, Type resultType,
                   Type elementType, Value basePtr, ValueRange indices,
-                  ArrayRef<NamedAttribute> attributes) {
+                  bool inbounds, ArrayRef<NamedAttribute> attributes) {
   build(builder, result, resultType, elementType, basePtr,
-        SmallVector<GEPArg>(indices), attributes);
+        SmallVector<GEPArg>(indices), inbounds, attributes);
 }
 
 static ParseResult
 parseGEPIndices(OpAsmParser &parser,
                 SmallVectorImpl<OpAsmParser::UnresolvedOperand> &indices,
                 DenseI32ArrayAttr &rawConstantIndices) {
   SmallVector<int32_t> constantIndices;
 

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

if (llvm::Function *callee = invokeInst->getCalledFunction()) {
           loc, types, operands, lookupBlock(invokeInst->getNormalDest()),
           normalArgs, lookupBlock(invokeInst->getUnwindDest()), unwindArgs);
     }
     if (!invokeInst->getType()->isVoidTy())
       mapValue(inst, invokeOp.getResults().front());
     return success();
   }
   if (inst->getOpcode() == llvm::Instruction::GetElementPtr) {
-    // FIXME: Support inbounds GEPs.
     auto *gepInst = cast<llvm::GetElementPtrInst>(inst);
     Type sourceElementType = convertType(gepInst->getSourceElementType());
     FailureOr<Value> basePtr = convertValue(gepInst->getOperand(0));
     if (failed(basePtr))
       return failure();
 
     // Treat every indices as dynamic since GEPOp::build will refine those
     // indices into static attributes later. One small downside of this
     // approach is that many unused `llvm.mlir.constant` would be emitted
     // at first place.
     SmallVector<GEPArg> indices;
     for (llvm::Value *operand : llvm::drop_begin(gepInst->operand_values())) {
       FailureOr<Value> index = convertValue(operand);
       if (failed(index))
         return failure();
       indices.push_back(index.value());
     }
 
     Type type = convertType(inst->getType());
-    Value res = builder.create<GEPOp>(loc, type, sourceElementType,
-                                      basePtr.value(), indices);
+    Value res =
+        builder.create<GEPOp>(loc, type, sourceElementType, basePtr.value(),
+                              indices, gepInst->isInBounds());
     mapValue(inst, res);
     return success();
   }
 
   return emitError(loc) << "unhandled instruction " << diag(*inst);
 }
 
 LogicalResult Importer::processInstruction(llvm::Instruction *inst) {
   // FIXME: Support uses of SubtargetData.
-  // FIXME: Add support for inbounds GEPs.
   // FIXME: Add support for fast-math flags and call / operand attributes.
   // FIXME: Add support for the indirectbr, cleanupret, catchret, catchswitch,
   // callbr, vaarg, landingpad, catchpad, cleanuppad instructions.
 
   // Convert LLVM intrinsics calls to MLIR intrinsics.
   if (auto *callInst = dyn_cast<llvm::CallInst>(inst)) {
     llvm::Function *callee = callInst->getCalledFunction();
     if (callee && callee->isIntrinsic())

mlir/test/Dialect/LLVMIR/canonicalize.mlir

 llvm.func @fold_gep(%x : !llvm.ptr<i8>) -> !llvm.ptr<i8> {
   %c0 = arith.constant 0 : i32
   %c = llvm.getelementptr %x[%c0] : (!llvm.ptr<i8>, i32) -> !llvm.ptr<i8>
   llvm.return %c : !llvm.ptr<i8>
 }
 
 // CHECK-LABEL: fold_gep_neg
 // CHECK-SAME: %[[a0:arg[0-9]+]]
-// CHECK-NEXT: %[[RES:.*]] = llvm.getelementptr %[[a0]][0, 1]
+// CHECK-NEXT: %[[RES:.*]] = llvm.getelementptr inbounds %[[a0]][0, 1]
 // CHECK-NEXT: llvm.return %[[RES]]
 llvm.func @fold_gep_neg(%x : !llvm.ptr) -> !llvm.ptr {
   %c0 = arith.constant 0 : i32
-  %0 = llvm.getelementptr %x[%c0, 1] : (!llvm.ptr, i32) -> !llvm.ptr, !llvm.struct<(i32, i32)>
+  %0 = llvm.getelementptr inbounds %x[%c0, 1] : (!llvm.ptr, i32) -> !llvm.ptr, !llvm.struct<(i32, i32)>
   llvm.return %0 : !llvm.ptr
 }
 
 // CHECK-LABEL: fold_gep_canon
 // CHECK-SAME: %[[a0:arg[0-9]+]]
 // CHECK-NEXT: %[[RES:.*]] = llvm.getelementptr %[[a0]][2]
 // CHECK-NEXT: llvm.return %[[RES]]
 llvm.func @fold_gep_canon(%x : !llvm.ptr<i8>) -> !llvm.ptr<i8> {

mlir/test/Dialect/LLVMIR/roundtrip.mlir

// CHECK:  llvm.return
   llvm.return
 }
 
 // CHECK-LABEL: @gep
 llvm.func @gep(%ptr: !llvm.ptr<struct<(i32, struct<(i32, f32)>)>>, %idx: i64,
                %ptr2: !llvm.ptr<struct<(array<10 x f32>)>>) {
   // CHECK: llvm.getelementptr %{{.*}}[%{{.*}}, 1, 0] : (!llvm.ptr<struct<(i32, struct<(i32, f32)>)>>, i64) -> !llvm.ptr<i32>
   llvm.getelementptr %ptr[%idx, 1, 0] : (!llvm.ptr<struct<(i32, struct<(i32, f32)>)>>, i64) -> !llvm.ptr<i32>
-  // CHECK: llvm.getelementptr %{{.*}}[%{{.*}}, 0, %{{.*}}] : (!llvm.ptr<struct<(array<10 x f32>)>>, i64, i64) -> !llvm.ptr<f32>
-  llvm.getelementptr %ptr2[%idx, 0, %idx] : (!llvm.ptr<struct<(array<10 x f32>)>>, i64, i64) -> !llvm.ptr<f32>
+  // CHECK: llvm.getelementptr inbounds %{{.*}}[%{{.*}}, 0, %{{.*}}] : (!llvm.ptr<struct<(array<10 x f32>)>>, i64, i64) -> !llvm.ptr<f32>
+  llvm.getelementptr inbounds %ptr2[%idx, 0, %idx] : (!llvm.ptr<struct<(array<10 x f32>)>>, i64, i64) -> !llvm.ptr<f32>
   llvm.return
 }
 
 // An larger self-contained function.
 // CHECK-LABEL: llvm.func @foo(%{{.*}}: i32) -> !llvm.struct<(i32, f64, i32)> {
 llvm.func @foo(%arg0: i32) -> !llvm.struct<(i32, f64, i32)> {
 // CHECK:  %[[V0:.*]] = llvm.mlir.constant(3 : i64) : i32
 // CHECK:  %[[V1:.*]] = llvm.mlir.constant(3 : i64) : i32

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

 }
 
 ; // -----
 
 ; CHECK-LABEL: @gep_static_idx
 ; CHECK-SAME:  %[[PTR:[a-zA-Z0-9]+]]
 define void @gep_static_idx(float* %ptr) {
   ; CHECK: %[[IDX:.+]] = llvm.mlir.constant(7 : i32)
-  ; CHECK: llvm.getelementptr %[[PTR]][%[[IDX]]] : (!llvm.ptr<f32>, i32) -> !llvm.ptr<f32>
-  %1 = getelementptr float, float* %ptr, i32 7
+  ; CHECK: llvm.getelementptr inbounds %[[PTR]][%[[IDX]]] : (!llvm.ptr<f32>, i32) -> !llvm.ptr<f32>
+  %1 = getelementptr inbounds float, float* %ptr, i32 7
   ret void
 }
 
 ; // -----
 
 %sub_struct = type { i32, i8 }
 %my_struct = type { %sub_struct, [4 x i32] }
 

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

define ptr @opaque_ptr_alloca(i32 %0) {
   %2 = alloca float, i32 %0, align 4
   ret ptr %2
 }
 
 ; // -----
 
 ; CHECK-LABEL: @opaque_ptr_gep
 define ptr @opaque_ptr_gep(ptr %0, i32 %1) {
-  ; CHECK: = llvm.getelementptr %{{.*}}[%{{.*}}] : (!llvm.ptr, i32) -> !llvm.ptr, f32
-  %3 = getelementptr float, ptr %0, i32 %1
+  ; CHECK: = llvm.getelementptr inbounds %{{.*}}[%{{.*}}] : (!llvm.ptr, i32) -> !llvm.ptr, f32
+  %3 = getelementptr inbounds float, ptr %0, i32 %1
   ret ptr %3
 }
 
 ; // -----
 
 ; CHECK-LABEL: @opaque_ptr_gep
 define ptr @opaque_ptr_gep_struct(ptr %0, i32 %1){
   ; CHECK: = llvm.getelementptr %{{.*}}[%{{.*}}, 0, 1] : (!llvm.ptr, i32) -> !llvm.ptr, !llvm.struct<(struct<(f32, f64)>, struct<(i32, i64)>)>
   %3 = getelementptr { { float, double }, { i32, i64 } }, ptr %0, i32 %1, i32 0, i32 1
   ret ptr %3
 }

mlir/test/Target/LLVMIR/llvmir.mlir

// CHECK-NEXT: fneg float %0
   llvm.return %10 : !llvm.struct<(f32, i32)>
 }
 
 // CHECK-LABEL: @gep
 llvm.func @gep(%ptr: !llvm.ptr<struct<(i32, struct<(i32, f32)>)>>, %idx: i64,
                %ptr2: !llvm.ptr<struct<(array<10 x f32>)>>) {
   // CHECK: = getelementptr { i32, { i32, float } }, ptr %{{.*}}, i64 %{{.*}}, i32 1, i32 0
   llvm.getelementptr %ptr[%idx, 1, 0] : (!llvm.ptr<struct<(i32, struct<(i32, f32)>)>>, i64) -> !llvm.ptr<i32>
-  // CHECK: = getelementptr { [10 x float] }, ptr %{{.*}}, i64 %{{.*}}, i32 0, i64 %{{.*}}
-  llvm.getelementptr %ptr2[%idx, 0, %idx] : (!llvm.ptr<struct<(array<10 x f32>)>>, i64, i64) -> !llvm.ptr<f32>
+  // CHECK: = getelementptr inbounds { [10 x float] }, ptr %{{.*}}, i64 %{{.*}}, i32 0, i64 %{{.*}}
+  llvm.getelementptr inbounds %ptr2[%idx, 0, %idx] : (!llvm.ptr<struct<(array<10 x f32>)>>, i64, i64) -> !llvm.ptr<f32>
   llvm.return
 }
 
 //
 // Indirect function calls
 //
 
 // CHECK-LABEL: define void @indirect_const_call(i64 {{%.*}})