[WebAssembly] WebAssemblyFastISel getelementptr variable index support

[jgravelle-google]

Summary

Previously -fast-isel getelementptr would constant-fold non-constant i8
load/stores.

[sunfish]

The underlying problem is that after it folds the dynamic index into the address base register, the code at the bottom of WebAssemblyFastISel::computeAddress just sticks the getelementptr base into the address base register, accidentally overwriting the previous value. A more precise fix would be to have that code at the bottom of WebAssemblyFastISel::computeAddress return false in the case where there's already a (non-zero) base register.

For extra caution, it may be possible to have Address::setReg assert that the base is zero before overwriting it.

[jgravelle-google]

• Reduce nesting in computeAddress, add test case for GEP flattening

[jgravelle-google]

• Rebase

[jgravelle-google]

• Remove commented out code

[sbc100]

lib/Target/WebAssembly/WebAssemblyFastISel.cpp

279

Is this continue now redundant?

[jgravelle-google]

In D34044#777275, @sunfish wrote:

A more precise fix would be to have that code at the bottom of WebAssemblyFastISel::computeAddress return false in the case where there's already a (non-zero) base register.

That's probably worth doing, the problem is that the line that does this problematic folding already checks Addr.getReg() == 0, so it doesn't actually fix the problem.
The real underlying cause is that we think we're folding Add instructions, but we haven't actually checked that the Use is an Add yet. Moving the unscaled add after the check makes that a correct assumption, and everything seems happy about it.
Note that adding the return false made the lit tests pass, and only failed on the wasm waterfall. I added a test case (@store_i8_with_array_alloca_gep) to cover it.
This also only fails with the -elf triple, because we store the stack pointer in memory instead of in a global, making allocas loads instead of get_globals.

For extra caution, it may be possible to have Address::setReg assert that the base is zero before overwriting it.

Probably a good idea going forward. Added.

@sunfish PTAL

[jgravelle-google]

• Remove unneeded continue

[jgravelle-google]

lib/Target/WebAssembly/WebAssemblyFastISel.cpp

279

Yep, thanks

[jgravelle-google]

• Re-add check before final setReg in computeAddress

[dschuff]

test/CodeGen/WebAssembly/offset-fastisel.ll

1

Add -fast-isel-abort=1

[jgravelle-google]

• Add fast-isel-abort=1 to test

[sunfish]

In D34044#779541, @jgravelle-google wrote:

In D34044#777275, @sunfish wrote:

A more precise fix would be to have that code at the bottom of WebAssemblyFastISel::computeAddress return false in the case where there's already a (non-zero) base register.

That's probably worth doing, the problem is that the line that does this problematic folding already checks Addr.getReg() == 0, so it doesn't actually fix the problem.
The real underlying cause is that we think we're folding Add instructions, but we haven't actually checked that the Use is an Add yet. Moving the unscaled add after the check makes that a correct assumption, and everything seems happy about it.

The if (S == 1 && Addr.isRegBase() && Addr.getReg() == 0) case doesn't depend on the use being an Add.

Note that adding the return false made the lit tests pass, and only failed on the wasm waterfall. I added a test case (@store_i8_with_array_alloca_gep) to cover it.

Ah, ok. Looking at that testcase, the problem is that the case Instruction::Alloca: code has another instance as the same bug as before; it also needs to return false if there's already a base register, before assigning a new base.

[jgravelle-google]

• Guard against Addr base being set in Alloca case

[jgravelle-google]

Remove unintended files

[jgravelle-google]

In D34044#782696, @sunfish wrote:

In D34044#779541, @jgravelle-google wrote:

In D34044#777275, @sunfish wrote:

A more precise fix would be to have that code at the bottom of WebAssemblyFastISel::computeAddress return false in the case where there's already a (non-zero) base register.

That's probably worth doing, the problem is that the line that does this problematic folding already checks Addr.getReg() == 0, so it doesn't actually fix the problem.
The real underlying cause is that we think we're folding Add instructions, but we haven't actually checked that the Use is an Add yet. Moving the unscaled add after the check makes that a correct assumption, and everything seems happy about it.

The if (S == 1 && Addr.isRegBase() && Addr.getReg() == 0) case doesn't depend on the use being an Add.

The comment // An unscaled add of a register. Set it as the new base. confused me then, especially because immediately following we are dealing with Adds.
Thinking about it more deeply, this is is basically when we have getelementptr i8, i8* %pointer, i32 %someExpression, and this lowers that to the wasm equivalent of (i32.add %pointer %someExpression)?

Note that adding the return false made the lit tests pass, and only failed on the wasm waterfall. I added a test case (@store_i8_with_array_alloca_gep) to cover it.

Ah, ok. Looking at that testcase, the problem is that the case Instruction::Alloca: code has another instance as the same bug as before; it also needs to return false if there's already a base register, before assigning a new base.

Ah, makes sense. Added asserts and an isSet function that should guard against these sorts of thing.

[jgravelle-google]

@sunfish , ping?

[sunfish]

Looks good!

lib/Target/WebAssembly/WebAssemblyFastISel.cpp

public:
     // Innocuous defaults for our address.
     Address() : Kind(RegBase), Offset(0), GV(0) { Base.Reg = 0; }
     void setKind(BaseKind K) { Kind = K; }
     BaseKind getKind() const { return Kind; }
     bool isRegBase() const { return Kind == RegBase; }
     bool isFIBase() const { return Kind == FrameIndexBase; }
     void setReg(unsigned Reg) {
       assert(isRegBase() && "Invalid base register access!");
+      assert(Base.Reg == 0 && "Overwriting non-zero register");
       Base.Reg = Reg;
     }
     unsigned getReg() const {
       assert(isRegBase() && "Invalid base register access!");
       return Base.Reg;
     }
     void setFI(unsigned FI) {
       assert(isFIBase() && "Invalid base frame index access!");
for (gep_type_iterator GTI = gep_type_begin(U), E = gep_type_end(U);
       } else {
         uint64_t S = DL.getTypeAllocSize(GTI.getIndexedType());
         for (;;) {
           if (const ConstantInt *CI = dyn_cast<ConstantInt>(Op)) {
             // Constant-offset addressing.
             TmpOffset += CI->getSExtValue() * S;
             break;
           }
+          if (!canFoldAddIntoGEP(U, Op)) {
+            goto unsupported_gep;
+          }
           if (S == 1 && Addr.isRegBase() && Addr.getReg() == 0) {
             // An unscaled add of a register. Set it as the new base.
             Addr.setReg(getRegForValue(Op));
             break;
           }
-          if (canFoldAddIntoGEP(U, Op)) {
-            // A compatible add with a constant operand. Fold the constant.
-            ConstantInt *CI =
-                cast<ConstantInt>(cast<AddOperator>(Op)->getOperand(1));
-            TmpOffset += CI->getSExtValue() * S;
-            // Iterate on the other operand.
-            Op = cast<AddOperator>(Op)->getOperand(0);
-            continue;
-          }
-          // Unsupported
-          goto unsupported_gep;
         }
       }
     }
     // Don't fold in negative offsets.
     if (int64_t(TmpOffset) >= 0) {
       // Try to grab the base operand now.
       Addr.setOffset(TmpOffset);
       if (computeAddress(U->getOperand(0), Addr))
         return true;
if (const ConstantInt *CI = dyn_cast<ConstantInt>(RHS)) {
       if (TmpOffset >= 0) {
         Addr.setOffset(TmpOffset);
         return computeAddress(LHS, Addr);
       }
     }
     break;
   }
   }
+  if (Addr.getReg() != 0) {
+    return false;
+  }
   Addr.setReg(getRegForValue(Obj));
   return Addr.getReg() != 0;
 }
 
 void WebAssemblyFastISel::materializeLoadStoreOperands(Address &Addr) {
   if (Addr.isRegBase()) {
     unsigned Reg = Addr.getReg();
     if (Reg == 0) {

test/CodeGen/WebAssembly/offset-fastisel.ll

+; RUN: llc < %s -asm-verbose=false -disable-wasm-explicit-locals -fast-isel | FileCheck %s
+
+; TODO: Merge this with offset.ll when fast-isel matches better.
+
+target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
+target triple = "wasm32-unknown-unknown-elf"
+
+; CHECK-LABEL: store_i8_with_variable_gep_offset:
+; CHECK: i32.add    $push[[L0:[0-9]+]]=, $0, $1{{$}}
+; CHECK: i32.const  $push[[L1:[0-9]+]]=, 0{{$}}
+; CHECK: i32.store8 0($pop[[L0]]), $pop[[L1]]{{$}}
+define void @store_i8_with_variable_gep_offset(i8* %p, i32 %idx) {
+  %s = getelementptr inbounds i8, i8* %p, i32 %idx
+  store i8 0, i8* %s
+  ret void
+}
+
+; CHECK-LABEL: store_i8_with_array_alloca_gep:
+; CHECK: i32.const   $push[[L0:[0-9]+]]=, 0{{$}}
+; CHECK: i32.load    $push[[L1:[0-9]+]]=, __stack_pointer($pop[[L0]]){{$}}
+; CHECK: i32.const   $push[[L2:[0-9]+]]=, 32{{$}}
+; CHECK: i32.sub     $push{{[0-9]+}}=, $pop[[L1]], $pop[[L2]]{{$}}
+; CHECK: i32.add     $push[[L4:[0-9]+]]=, $pop{{[0-9]+}}, $0{{$}}
+; CHECK: i32.const   $push[[L5:[0-9]+]]=, 0{{$}}
+; CHECK: i32.store8  0($pop[[L4]]), $pop[[L5]]{{$}}
+define hidden void @store_i8_with_array_alloca_gep(i32 %idx) {
+  %A = alloca [30 x i8], align 16
+  %s = getelementptr inbounds [30 x i8], [30 x i8]* %A, i32 0, i32 %idx
+  store i8 0, i8* %s, align 1
+  ret void
+}
+
+; CHECK-LABEL: store_i32_with_unfolded_gep_offset:
+; CHECK: i32.const $push[[L0:[0-9]+]]=, 24{{$}}
+; CHECK: i32.add   $push[[L1:[0-9]+]]=, $0, $pop[[L0]]{{$}}
+; CHECK: i32.const $push[[L2:[0-9]+]]=, 0{{$}}
+; CHECK: i32.store 0($pop[[L1]]), $pop[[L2]]{{$}}
+define void @store_i32_with_unfolded_gep_offset(i32* %p) {
+  %s = getelementptr i32, i32* %p, i32 6
+  store i32 0, i32* %s
+  ret void
+}
+
+; CHECK-LABEL: store_i32_with_folded_gep_offset:
+; CHECK: i32.store 24($0), $pop{{[0-9]+$}}
+define void @store_i32_with_folded_gep_offset(i32* %p) {
+  %s = getelementptr inbounds i32, i32* %p, i32 6
+  store i32 0, i32* %s
+  ret void
+}
+
+; CHECK-LABEL: load_i32_with_folded_gep_offset:
+; CHECK: i32.load  $push{{[0-9]+}}=, 24($0){{$}}
+define i32 @load_i32_with_folded_gep_offset(i32* %p) {
+  %s = getelementptr inbounds i32, i32* %p, i32 6
+  %t = load i32, i32* %s
+  ret i32 %t
+}
+
+; CHECK-LABEL: store_i64_with_unfolded_gep_offset:
+; CHECK: i32.const $push[[L0:[0-9]+]]=, 24{{$}}
+; CHECK: i32.add   $push[[L1:[0-9]+]]=, $0, $pop[[L0]]{{$}}
+; CHECK: i64.const $push[[L2:[0-9]+]]=, 0{{$}}
+; CHECK: i64.store 0($pop[[L1]]), $pop[[L2]]{{$}}
+define void @store_i64_with_unfolded_gep_offset(i64* %p) {
+  %s = getelementptr i64, i64* %p, i32 3
+  store i64 0, i64* %s
+  ret void
+}
+
+; CHECK-LABEL: store_i8_with_folded_gep_offset:
+; CHECK: i32.store8 24($0), $pop{{[0-9]+$}}
+define void @store_i8_with_folded_gep_offset(i8* %p) {
+  %s = getelementptr inbounds i8, i8* %p, i32 24
+  store i8 0, i8* %s
+  ret void
+}
+
+; CHECK-LABEL: load_i8_u_with_folded_offset:
+; CHECK: i32.load8_u $push{{[0-9]+}}=, 24($0){{$}}
+define i32 @load_i8_u_with_folded_offset(i8* %p) {
+  %q = ptrtoint i8* %p to i32
+  %r = add nuw i32 %q, 24
+  %s = inttoptr i32 %r to i8*
+  %t = load i8, i8* %s
+  %u = zext i8 %t to i32
+  ret i32 %u
+}
+
+; TODO: this should be load8_s, need to fold sign-/zero-extend in fast-isel
+; CHECK-LABEL: load_i8_s_with_folded_offset:
+; CHECK: i32.load8_u $push{{[0-9]+}}=, 24($0){{$}}
+define i32 @load_i8_s_with_folded_offset(i8* %p) {
+  %q = ptrtoint i8* %p to i32
+  %r = add nuw i32 %q, 24
+  %s = inttoptr i32 %r to i8*
+  %t = load i8, i8* %s
+  %u = sext i8 %t to i32
+  ret i32 %u
+}