[MLIR][crashfix] Fall back to malloc when alignment is 0 for align-alloc

[kernhanda]

Summary

Currently, the following crashes when running the StandardToLLVM conversion
pass with the use-aligned-alloc option.

module {
  func @f() {
    %0 = alloc() {alignment = 0} : memref<64xf32>
    return
  }
}

To fix this, I chose the option falling back to malloc instead of emitting
an error. This is because aligned_alloc is not supposed to be called with 0
as a given alignment, so this allows for an escape hatch in the conversion,
while also fixing the crash.

[kernhanda]

clang-format

[bondhugula]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1524

The alignment is actually required to be a power of two for aligned_alloc. The code above was written with the assumption that the alignment would be a power of two. Unfortunately, we don't document what a valid alignment on the AllocOp is. Unlike llvm's alloca where a 0 alignment is treated the same as not specifying an alignment, we could make the alloc op's verifier fail for non power of two alignments. This will fix this case as well - since zero is not a power of two.

This fix will silently fall back to malloc for '0' alignment instead of trying to use aligned_alloc with elt size alignment (which is what happens if no alignment is set). If we wanted to handle this case, the right fix would be to check for a non-power of two alignment (in getAllocationAlignment), and use the default elt size based alignment while still using aligned_alloc.

[kernhanda]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1524

I'm okay with either of the alternatives, if others have an opinion.

[ftynse]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1524

I agree with @bondhugula that it should be fixed at the op verifier level. Arguably, silently falling back to the strategy that was explicitly disabled by an option (i.e. using malloc instead of aligned_alloc) is a surprising and undesirable behavior.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

if (auto allocaOp = dyn_cast<AllocaOp>(op)) {
                                : 0);
     }
 
     // Heap allocations.
     AllocOp allocOp = cast<AllocOp>(op);
 
     Optional<int64_t> allocationAlignment = getAllocationAlignment(allocOp);
     // Whether to use std lib function aligned_alloc that supports alignment.
-    bool useAlignedAlloc = allocationAlignment.hasValue();
+    bool useAlignedAlloc = allocationAlignment && *allocationAlignment != 0;
 
     // Insert the malloc/aligned_alloc declaration if it is not already present.
     auto allocFuncName = useAlignedAlloc ? "aligned_alloc" : "malloc";
     auto module = allocOp.getParentOfType<ModuleOp>();
     auto allocFunc = module.lookupSymbol<LLVM::LLVMFuncOp>(allocFuncName);
     if (!allocFunc) {
       OpBuilder moduleBuilder(op->getParentOfType<ModuleOp>().getBodyRegion());
       SmallVector<LLVM::LLVMType, 2> callArgTypes = {getIndexType()};

mlir/test/Conversion/StandardToLLVM/convert-dynamic-memref-ops.mlir

// ALIGNED-ALLOC-NEXT:  llvm.bitcast %[[allocated]] : !llvm<"i8*"> to !llvm<"float*">
   // ALIGNED-ALLOC-NEXT:  llvm.urem
   // ALIGNED-ALLOC-NEXT:  %[[SIZE_ALIGNED:.*]] = llvm.add
   // ALIGNED-ALLOC-NEXT:  llvm.call @aligned_alloc(%[[c32]], %[[SIZE_ALIGNED]])
   %5 = alloc() {alignment = 32} : memref<100xf32>
   // Bump alignment to the next power of two if it isn't.
   // ALIGNED-ALLOC:  %[[c128:.*]] = llvm.mlir.constant(128 : i64) : !llvm.i64
   // ALIGNED-ALLOC:  llvm.call @aligned_alloc(%[[c128]]
   %6 = alloc(%N) : memref<?xvector<18xf32>>
+  // When alignment is 0, revert to malloc behavior
+  // ALIGNED-ALLOC:  %[[c64:.*]] = llvm.mlir.constant(64 : index) : !llvm.i64
+  // ALIGNED-ALLOC:  llvm.mul %[[c64]]
+  // ALIGNED-ALLOC:  llvm.add
+  // ALIGNED-ALLOC-NEXT:  %[[ZERO_ALIGNED:.*]] = llvm.sub
+  // ALIGNED-ALLOC-NEXT:  llvm.call @malloc(%[[ZERO_ALIGNED]])
+  %7 = alloc() {alignment = 0} : memref<64xf32>
   return %0 : memref<32x18xf32>
 }
 
 // CHECK-LABEL: func @mixed_load(
 // CHECK-COUNT-2: !llvm<"float*">,
 // CHECK-COUNT-5: {{%[a-zA-Z0-9]*}}: !llvm.i64
 // CHECK:         %[[I:.*]]: !llvm.i64,
 // CHECK:         %[[J:.*]]: !llvm.i64)