Sign- and zero-extend 8- and 16-bit arguments to declared functions

tslang/lib/TypeScript/LowerToLLVM.cpp

@@ -1492,6 +1492,31 @@ struct ReturnInternalOpLowering : public TsLlvmPattern<mlir_ts::ReturnInternalOp
     }
 };
 
+// `llvm.signext` for s8/s16, `llvm.zeroext` for u8/u16 and boolean; empty for anything else,
+// including signless i8/i16, whose extension no type says.
+static StringRef getIntegerExtensionAttrName(mlir::Type type)
+{
+    if (isa<mlir_ts::BooleanType>(type))
+    {
+        return LLVM::LLVMDialect::getZExtAttrName();
+    }
+
+    if (auto intType = dyn_cast<mlir::IntegerType>(type); intType && intType.getWidth() < 32)
+    {
+        if (intType.isSigned())
+        {
+            return LLVM::LLVMDialect::getSExtAttrName();
+        }
+
+        if (intType.isUnsigned())
+        {
+            return LLVM::LLVMDialect::getZExtAttrName();
+        }
+    }
+
+    return {};
+}
+
 struct FuncOpLowering : public TsLlvmPattern<mlir_ts::FuncOp>
 {
     using TsLlvmPattern<mlir_ts::FuncOp>::TsLlvmPattern;
@@ -1528,6 +1553,30 @@ struct FuncOpLowering : public TsLlvmPattern<mlir_ts::FuncOp>
         auto convertedFuncType = rewriter.getFunctionType(signatureInputsConverter.getConvertedTypes(), signatureResultsConverter.getConvertedTypes());
         auto newFuncOp = rewriter.create<mlir::func::FuncOp>(location, funcOp.getName(), convertedFuncType, ArrayRef<NamedAttribute>{}, argDictAttrs);
 
+        // C expects an 8- or 16-bit integer or a bool extended to 32 bits by whoever hands it
+        // over: the caller for an argument, the callee for a result, and clang-compiled code
+        // relies on it (x86-64 SysV). A declaration - implemented elsewhere, usually in C - says
+        // so. A definition does not: a `signext` parameter makes the callee trust its caller, and
+        // tslang's indirect calls carry no call-site attributes.
+        if (funcOp.getBody().empty())
+        {
+            for (auto [index, type] : llvm::enumerate(fnType.getInputs()))
+            {
+                if (auto extension = getIntegerExtensionAttrName(type); !extension.empty())
+                {
+                    newFuncOp.setArgAttr(index, extension, rewriter.getUnitAttr());
+                }
+            }
+
+            for (auto [index, type] : llvm::enumerate(fnType.getResults()))
+            {
+                if (auto extension = getIntegerExtensionAttrName(type); !extension.empty())
+                {
+                    newFuncOp.setResultAttr(index, extension, rewriter.getUnitAttr());
+                }
+            }
+        }
+
         for (const auto &namedAttr : funcOp->getAttrs())
         {
             if (namedAttr.getName() == funcOp.getFunctionTypeAttrName())

tslang/test/tester/CMakeLists.txt

@@ -2305,6 +2305,29 @@ set_tests_properties(test-compile-linkname-declare-not-exported
                      PROPERTIES PASS_REGULAR_EXPRESSION "declare i64 @strlen\\(ptr\\)"
                                 FAIL_REGULAR_EXPRESSION "declare dllexport|Stack dump|Assertion failed")
 
+# C code expects an 8- or 16-bit integer (and a bool) extended to 32 bits by whoever hands it
+# over: the caller for an argument, the callee for a result. On x86-64 SysV clang relies on it,
+# so without `signext`/`zeroext` a call like take_s8(-5) passed garbage upper bits (tsbindgen's
+# `narrow` row read 251). A `declare`d function gets the attributes; `i32` is left alone.
+set(small_int_ext_names s8 s16 u8 u16 bool s32 result_s8 result_u16)
+set(small_int_ext_patterns
+    "declare i32 @take_s8\\(i8 signext\\)"
+    "declare i32 @take_s16\\(i16 signext\\)"
+    "declare i32 @take_u8\\(i8 zeroext\\)"
+    "declare i32 @take_u16\\(i16 zeroext\\)"
+    "declare i32 @take_bool\\(i1 zeroext\\)"
+    "declare i32 @take_s32\\(i32\\)"
+    "declare signext i8 @give_s8\\(\\)"
+    "declare zeroext i16 @give_u16\\(\\)")
+foreach(small_int_ext_name small_int_ext_pattern IN ZIP_LISTS small_int_ext_names small_int_ext_patterns)
+    add_test(NAME test-compile-small-int-ext-${small_int_ext_name}
+             COMMAND $<TARGET_FILE:tslang> --emit=llvm --no-default-lib -mm=none
+                     "${PROJECT_SOURCE_DIR}/test/tester/small-int-ext/declare.ts" -o -)
+    set_tests_properties(test-compile-small-int-ext-${small_int_ext_name}
+                         PROPERTIES PASS_REGULAR_EXPRESSION "${small_int_ext_pattern}"
+                                    FAIL_REGULAR_EXPRESSION "Stack dump|Assertion failed")
+endforeach()
+
 # A statement's condition that produces no value is an error, not a statement generated
 # without its test (an `if` that vanished, an endless `for`, a segfault in `do`, an "empty
 # block" error from `while`). The void call

tslang/test/tester/small-int-ext/declare.ts

@@ -0,0 +1,15 @@
+// C expects 8- and 16-bit integers extended to 32 bits by whoever hands them over: the caller
+// for an argument, the callee for a result (clang relies on both on x86-64 SysV).
+
+declare function take_s8(v: s8): s32;
+declare function take_s16(v: s16): s32;
+declare function take_u8(v: u8): s32;
+declare function take_u16(v: u16): s32;
+declare function take_bool(v: boolean): s32;
+declare function take_s32(v: s32): s32;
+declare function give_s8(): s8;
+declare function give_u16(): u16;
+
+function main() {
+    print(take_s8(-5), take_s16(-3), take_u8(255), take_u16(65535), take_bool(true), take_s32(1), give_s8(), give_u16());
+}