Sign extension not properly propagated to users

Reproduce the failure with the supplied cgeist command and sqrt.cc, then inspect the generated MLIR around the arith.extsi and arith.subi operations. Trace how the sign-extended value is used in the subtraction and update the frontend so the operation has consistent operand types. Rerun the command and confirm the generated IR verifies without the type error.

Given the following C++ code:

short isqrt(short num) {
  short res = 0;
  short bit = 1 << 14;  // ((unsigned) INT16_MAX + 1) / 2.

  for (int i = 0; i < 8; ++i) {
    if (num >= res + bit) {
      num -= res + bit;
      res = (res >> 1) + bit;
    } else {
      res >>= 1;
    }
    bit >>= 2;
  }
  return res;
}

Running cgeist at 4b04755a63fc as follows:

bin/cgeist \
  '-function=*' \
  -raise-scf-to-affine \
  --memref-fullrank \
  -S \
  -O0 \
  sqrt.cc

Produces

loc("/path/to/sqrt.cc":25:11): error: 'arith.subi' op requires the same type for all operands and results
"builtin.module"() ({
  "func.func"() <{function_type = (i16) -> i16, sym_name = "_Z5isqrts"}> ({
  ^bb0(%arg0: i16):
    %0 = "arith.constant"() <{value = 16384 : i16}> : () -> i16
    %1 = "arith.constant"() <{value = 2 : i16}> : () -> i16
    %2 = "arith.constant"() <{value = 1 : i16}> : () -> i16
    %3 = "arith.constant"() <{value = 1 : i32}> : () -> i32
    %4 = "arith.constant"() <{value = 0 : i16}> : () -> i16
    %5:3 = "affine.for"(%0, %4, %arg0) ({
    ^bb0(%arg1: index, %arg2: i16, %arg3: i16, %arg4: i16):
      %6 = "arith.extsi"(%arg4) : (i16) -> i32
      %7 = "arith.extsi"(%arg3) : (i16) -> i32
      %8 = "arith.extsi"(%arg2) : (i16) -> i32
      %9 = "arith.addi"(%7, %8) : (i32, i32) -> i32
      %10 = "arith.cmpi"(%6, %9) <{predicate = 5 : i64}> : (i32, i32) -> i1
      %11:2 = "scf.if"(%10) ({
        %13 = "arith.subi"(%arg4, %9) : (i16, i32) -> i16
        %14 = "arith.shrsi"(%7, %3) : (i32, i32) -> i32
        %15 = "arith.addi"(%14, %8) : (i32, i32) -> i32
        %16 = "arith.trunci"(%15) : (i32) -> i16
        "scf.yield"(%16, %13) : (i16, i16) -> ()
      }, {
        %13 = "arith.shrsi"(%arg3, %2) : (i16, i16) -> i16
        "scf.yield"(%13, %arg4) : (i16, i16) -> ()
      }) : (i1) -> (i16, i16)
      %12 = "arith.shrsi"(%arg2, %1) : (i16, i16) -> i16
      "affine.yield"(%12, %11#0, %11#1) : (i16, i16, i16) -> ()
    }) {lower_bound = affine_map<() -> (0)>, step = 1 : index, upper_bound = affine_map<() -> (8)>} : (i16, i16, i16) -> (i16, i16, i16)
    "func.return"(%5#1) : (i16) -> ()
  }) {llvm.linkage = #llvm.linkage<external>} : () -> ()
}) {<lots of attributes>} : () -> ()

It looks like the issue is that

%13 = "arith.subi"(%arg4, %9) : (i16, i32) -> i16

Is still using %arg4 even though it was sign extended in an earlier line

%6 = "arith.extsi"(%arg4) : (i16) -> i32