fix(rocke): emit magic division's sum without signed-overflow poison

dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml

@@ -35,6 +35,7 @@ test_categories:
       - "rocke_future_intrinsic_lowering"
       - "rocke_ir_types"
       - "rocke_opcode_values"
+      - "rocke_unsigned_int_ops"
       - "rocke_tiled_attention_2d_reentrancy"
       - "rocke_unified_attention_scalar_lower_to_llvm"
       - "rocke_python_import_smoke"
@@ -71,6 +72,7 @@ test_categories:
       - "rocke_future_intrinsic_lowering"
       - "rocke_ir_types"
       - "rocke_opcode_values"
+      - "rocke_unsigned_int_ops"
       - "rocke_tiled_attention_2d_reentrancy"
       - "rocke_unified_attention_scalar_lower_to_llvm"
       - "rocke_python_import_smoke"

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_conv_implicit_gemm_common.py

@@ -39,11 +39,11 @@
     Type,
     Value,
 )
+from rocke.helpers.magic_division import do_magic_division_dynamic
 from rocke.helpers.spec import choose_load_vec
 from rocke.helpers.transforms import (
     TensorDescriptor,
     DynamicTensorDescriptor,
-    do_magic_division_dynamic,
     embed,
     embed_dynamic,
     pad,

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py

@@ -53,10 +53,9 @@
     use_2d_kernel,
 )
 
+from rocke.helpers.magic_division import calculate_magic_numbers, do_magic_division
 from rocke.helpers.transforms import (
     TensorDescriptor,
-    calculate_magic_numbers,
-    do_magic_division,
 )
 
 # Alias the promoted log2-domain softcap helper under the local name used by
@@ -244,7 +243,7 @@ def _magic_div(b: IRBuilder, dividend: Value, divisor: int) -> Value:
     (a grid axis or a KV position), well inside the documented 31-bit
     unsigned range of the magic sequence -- so the unsigned magic quotient
     equals the signed ``b.div`` it replaces. This is the device lowering of
-    ``merge_v2_magic_division`` (``transforms.do_magic_division``).
+    ``merge_v2_magic_division`` (``magic_division.do_magic_division``).
     """
     mult, shift = calculate_magic_numbers(divisor)
     return do_magic_division(b, dividend, mult, shift)

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_args.py

@@ -32,8 +32,8 @@
     conv_direct_arg_names,
 )
 from rocke.helpers.spec import ceil_div
+from rocke.helpers.magic_division import calculate_magic_numbers
 from rocke.helpers.streamk import StreamKIterPartition, streamk_iter_partition
-from rocke.helpers.transforms import calculate_magic_numbers
 
 __all__ = [
     "ConvArgs",

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_dgrad.py

@@ -69,15 +69,17 @@
 from rocke.helpers.geometry import WarpGrid
 from rocke.helpers.layouts import ConvKOuterFragmentReader, LdsLayout
 from rocke.helpers.loads import CoalescedTileLoader
+from rocke.helpers.magic_division import (
+    calculate_magic_numbers,
+    do_magic_division_dynamic,
+)
 from rocke.helpers.mfma_gemm_inner import decode_mfma_lanes
 from rocke.helpers.schedule import SchedulePolicy
 from rocke.helpers.spec import ceil_div, kernel_name_join
 from rocke.helpers.tensor_view import make_buffer_resource
 from rocke.helpers.transforms import (
     DynamicTensorDescriptor,
     TensorDescriptor,
-    calculate_magic_numbers,
-    do_magic_division_dynamic,
     embed,
     pad,
     unmerge_magic,

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py

@@ -2501,7 +2501,7 @@ def _free_axis_vec(chan: int, dtype: str, override: "int | None" = None) -> int:
             # so that run is exactly Gm, a build-time constant: fold it and its
             # magic pair here instead of reading the kernargs. The epilogue
             # keeps the true cpg through ``params``.
-            from rocke.helpers.transforms import calculate_magic_numbers
+            from rocke.helpers.magic_division import calculate_magic_numbers
 
             _merged_cpg = p_load.cpg
             _mult, _shift = calculate_magic_numbers(_merged_cpg)

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/fmha_splitkv_decode.py

@@ -40,8 +40,8 @@
     store_scalar_from_f32,
     store_vec,
 )
+from rocke.helpers.magic_division import calculate_magic_numbers, do_magic_division
 from rocke.helpers.spec import kernel_name_join
-from rocke.helpers.transforms import calculate_magic_numbers, do_magic_division
 from ._fmha_common import FmhaCommonSpec, FmhaKernelBuilder, validate_common_spec
 from ._fmha_warp_body import WARP_SIZE
 from .fmha_arch import validate_fmha_mfma_atom

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/sage_attention.py

@@ -67,6 +67,7 @@
 )
 from rocke.helpers.i4_dequant import unpack_i4_byte_to_pair_i32
 from rocke.helpers.io import io_ir_type, load_scalar_as_f32
+from rocke.helpers.magic_division import calculate_magic_numbers, do_magic_division
 from rocke.helpers.mfma_attention import (
     MFMA_ATTN_BLOCK_K,
     MFMA_ATTN_BLOCK_M,
@@ -79,7 +80,6 @@
     load_q_scale_for_block,
 )
 from rocke.helpers.spec import kernel_name_join
-from rocke.helpers.transforms import calculate_magic_numbers, do_magic_division
 from ._fmha_common import FmhaCommonSpec, FmhaKernelBuilder, validate_common_spec
 from ._fmha_warp_body import WARP_SIZE, fmha_warp_fwd_inner_body
 from .fmha_arch import validate_fmha_mfma_atom

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/sparse_attention.py

@@ -46,13 +46,13 @@
 from typing import Tuple
 
 from rocke.core.ir import I8, IRBuilder, KernelDef, Value
+from rocke.helpers.magic_division import calculate_magic_numbers, do_magic_division
 from rocke.helpers.mfma_attention import (
     MFMA_ATTN_BLOCK_K,
     MFMA_ATTN_BLOCK_M,
     mfma_attention_fwd_inner_body,
 )
 from rocke.helpers.spec import kernel_name_join
-from rocke.helpers.transforms import calculate_magic_numbers, do_magic_division
 from ._fmha_common import FmhaCommonSpec, FmhaKernelBuilder, validate_common_spec
 from .fmha_arch import validate_fmha_mfma_atom
 

dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py

@@ -2345,15 +2345,15 @@ def attention_dense_signature(spec: AttentionDenseSpec):
 
 def _fastdiv_kernargs(spec: Gfx942AttentionDenseSpec) -> dict[str, int]:
     """``{<divisor>_magic, <divisor>_shift}`` for every ``spec.fastdiv_divisors``
-    entry, from :func:`rocke.helpers.transforms.calculate_magic_numbers`. The magic
+    entry, from :func:`rocke.helpers.magic_division.calculate_magic_numbers`. The magic
     is a uint32 passed as its two's-complement i32 bit pattern, which is what the
     kernel's unsigned mul-hi reads.
 
     Raises when the work-item count reaches the fast-division dividend limit: every
     decode dividend is below it, and past it the quotient is wrong with no fault."""
     if not spec.fastdiv_divisors:
         return {}
-    from rocke.helpers.transforms import calculate_magic_numbers
+    from rocke.helpers.magic_division import calculate_magic_numbers
 
     nqb = spec.seqlen_q // spec.block_m
     work = nqb * spec.num_query_heads * spec.batch

dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx942_dense_wiring.py

@@ -46,7 +46,7 @@
     build_attention_dense,
     supports_attention_dense,
 )
-from rocke.helpers.transforms import calculate_magic_numbers
+from rocke.helpers.magic_division import calculate_magic_numbers
 
 _NAME = "attention_gfx942_dense"
 _SPEC_ID = "gfx942_dense"

dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py

@@ -4220,7 +4220,7 @@ def test_persistent_kernargs_append_fastdiv_pairs_per_decode(self):
         divisor of the resolved decode, in the built kernel's param order; the
         launcher's values are ``calculate_magic_numbers`` of the divisors, packed
         as two's-complement i32, and divide every work-item index exactly."""
-        from rocke.helpers.transforms import calculate_magic_numbers
+        from rocke.helpers.magic_division import calculate_magic_numbers
         from kernels.gfx942.attention_dense import (
             attention_dense_runtime_args,
             attention_dense_signature,

dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py

@@ -2108,13 +2108,13 @@ def test_persistent_kernarg_abi_carries_the_fastdiv_pairs(decode):
     GPU only.
 
     The values are checked two ways: equal to
-    :func:`rocke.helpers.transforms.calculate_magic_numbers` for the divisor (magic
+    :func:`rocke.helpers.magic_division.calculate_magic_numbers` for the divisor (magic
     packed as its two's-complement i32), and -- through the host model of the
     in-kernel divide -- exact for dividends up to just under the 2**30 bound. Non-pow2
     divisors throughout (B=3, NQB=13, Hq=40, gqa=5), so the magic is not the trivial
     pow2 one.
     """
-    from rocke.helpers.transforms import calculate_magic_numbers
+    from rocke.helpers.magic_division import calculate_magic_numbers
 
     spec = _persistent_spec(
         batch=3,

dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt

@@ -415,6 +415,17 @@ if(ROCKE_INSTALL_TESTS)
            "set_tests_properties(rocke_golden_static PROPERTIES LABELS \"ckc;python;golden;host;quick\")\n")
 
     if(ROCKE_INSTALL_PLATFORM_PYTEST)
+      # Native helpers some Python tests drive, as paths relative to the
+      # CTestTestfile: the portable-IR replay CLI, and the C API program whose
+      # --print-llvm output test_unsigned_int_ops.py compares with the Python
+      # builder's text.
+      set(_rocke_pytest_env
+          "PYTHONPATH=.;ROCKE_REPLAY_CLI=${ROCKE_TEST_CTEST_REL}tests/portable_ir/rocke_portable_ir_replay_cli${CMAKE_EXECUTABLE_SUFFIX}")
+      if("rocke_unsigned_int_ops" IN_LIST _ROCKE_INSTALL_TEST_TARGETS)
+        rocke_cxx_test_installed_name(_rocke_installed_name rocke_unsigned_int_ops)
+        string(APPEND _rocke_pytest_env
+               ";ROCKE_UNSIGNED_INT_OPS_CLI=${ROCKE_CXX_TEST_CTEST_REL}${_rocke_installed_name}${CMAKE_EXECUTABLE_SUFFIX}")
+      endif()
       # Full pytest suite. PYTHONPATH points at the CTestTestfile's working
       # directory (which contains the installed rocke package) so `import rocke`
       # resolves without the source-tree sys.path shims. The library tests have
@@ -424,7 +435,7 @@ if(ROCKE_INSTALL_TESTS)
       string(APPEND _ROCKE_INSTALLED_CTEST_ENTRIES
              "set_tests_properties(rocke_pytest PROPERTIES\n"
              "  LABELS \"ckc\\;python\\;quick\"\n"
-             "  ENVIRONMENT \"PYTHONPATH=.;ROCKE_REPLAY_CLI=${ROCKE_TEST_CTEST_REL}tests/portable_ir/rocke_portable_ir_replay_cli${CMAKE_EXECUTABLE_SUFFIX}\"\n"
+             "  ENVIRONMENT \"${_rocke_pytest_env}\"\n"
              "  TIMEOUT 300)\n")
     endif()
 

dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp

@@ -519,6 +519,88 @@ const char* rocke_attr_get_str(const rocke_attr_map_t* m, const char* key)
     return v->u.s;
 }
 
+/* ir._INT_ADD_WRAPS / ir._INT_CMP_PREDICATES. */
+static const char* const k_int_add_wraps[] = {"nsw", "nuw", NULL};
+static const char* const k_int_cmp_predicates[] = {"lt", "le", "gt", "ge", "eq", "ne", "ult", NULL};
+
+/* The shared body of rocke_i_add_wrap_error / rocke_i_cmp_pred_error: the
+ * error for the string attr `key` of `op_name`, or NULL when it is absent or
+ * one of `allowed`. The unsigned value `unsigned_value` takes i32 or i64
+ * operands only. */
+static const char* string_attr_error(const rocke_attr_map_t* attrs,
+                                     const char* op_name,
+                                     const char* key,
+                                     const char* const* allowed,
+                                     const char* not_a_string,
+                                     const char* unsigned_value,
+                                     const rocke_type_t* operand_type,
+                                     char* buf,
+                                     size_t cap)
+{
+    const rocke_attr_value_t* v = rocke_attr_get(attrs, key);
+    if(!v)
+        return NULL;
+    if(v->kind != ROCKE_ATTR_STR)
+        return not_a_string;
+    bool known = false;
+    for(const char* const* s = allowed; *s; ++s)
+    {
+        if(strcmp(v->u.s, *s) == 0)
+            known = true;
+    }
+    if(!known)
+    {
+        snprintf(buf, cap, "%s: unknown %s '%s'", op_name, key, v->u.s);
+        return buf;
+    }
+    const char* type_name = operand_type ? operand_type->name : "?";
+    if(strcmp(v->u.s, unsigned_value) == 0 && strcmp(type_name, "i32") != 0
+       && strcmp(type_name, "i64") != 0)
+    {
+        snprintf(buf,
+                 cap,
+                 "%s: %s '%s' needs i32 or i64 operands, got %s",
+                 op_name,
+                 key,
+                 unsigned_value,
+                 type_name);
+        return buf;
+    }
+    return NULL;
+}
+
+const char* rocke_i_add_wrap_error(const rocke_attr_map_t* attrs,
+                                   const rocke_type_t* operand_type,
+                                   char* buf,
+                                   size_t cap)
+{
+    return string_attr_error(attrs,
+                             "arith.add",
+                             "wrap",
+                             k_int_add_wraps,
+                             "arith.add: wrap must be the string 'nsw' or 'nuw'",
+                             "nuw",
+                             operand_type,
+                             buf,
+                             cap);
+}
+
+const char* rocke_i_cmp_pred_error(const rocke_attr_map_t* attrs,
+                                   const rocke_type_t* operand_type,
+                                   char* buf,
+                                   size_t cap)
+{
+    return string_attr_error(attrs,
+                             "arith.cmp",
+                             "pred",
+                             k_int_cmp_predicates,
+                             "arith.cmp: pred must be a string",
+                             "ult",
+                             operand_type,
+                             buf,
+                             cap);
+}
+
 bool rocke_attr_get_bool(const rocke_attr_map_t* m, const char* key, bool dflt)
 {
     const rocke_attr_value_t* v = rocke_attr_get(m, key);
@@ -804,7 +886,10 @@ static const char* const rocke_opcode_names[ROCKE_OP__COUNT] = {
     "arith.umax",
     "arith.umin",
     "arith.uitofp_f32",
-    "arith.cvt_f32_to_u8_sat"};
+    "arith.cvt_f32_to_u8_sat",
+
+    /* arith.* range marker */
+    "arith.assume_range"};
 
 const char* rocke_opcode_name(rocke_opcode_t op)
 {
@@ -1083,7 +1168,10 @@ static const bool rocke_opcode_pure[ROCKE_OP__COUNT] = {
     /* arith.umax                 */ true,
     /* arith.umin                 */ true,
     /* arith.uitofp_f32           */ true,
-    /* arith.cvt_f32_to_u8_sat    */ true};
+    /* arith.cvt_f32_to_u8_sat    */ true,
+
+    /* arith.* range marker */
+    /* arith.assume_range         */ true};
 
 bool rocke_opcode_is_pure(rocke_opcode_t op)
 {

dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp

@@ -146,6 +146,48 @@ rocke_value_t* rocke_b_add(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_
     return rocke_i_binop(b, ROCKE_OP_ARITH_ADD, a, c, "add");
 }
 
+/* Python _check_unsigned_int_operands: both lowerings read the operands as
+ * uint32_t / uint64_t by width, so they are two i32 or two i64 values. Returns
+ * false (with the builder error set) when they are not. */
+static bool check_unsigned_int_operands(rocke_ir_builder_t* b,
+                                        const char* op_name,
+                                        const rocke_value_t* a,
+                                        const rocke_value_t* c)
+{
+    if(!a || !c)
+    {
+        rocke_i_set_err(b, ROCKE_ERR_VALUE, "%s NULL operand", op_name);
+        return false;
+    }
+    if(!rocke_type_eq(a->type, c->type)
+       || (!rocke_i_type_is(a->type, "i32") && !rocke_i_type_is(a->type, "i64")))
+    {
+        rocke_i_set_err(b,
+                        ROCKE_ERR_VALUE,
+                        "%s expects two i32 or two i64 operands, got %s / %s",
+                        op_name,
+                        a->type->name,
+                        c->type->name);
+        return false;
+    }
+    return true;
+}
+
+rocke_value_t* rocke_b_add_nuw(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c)
+{
+    rocke_value_t* operands[2];
+    rocke_attr_map_t at;
+    if(!rocke_i_live(b))
+        return NULL;
+    if(!check_unsigned_int_operands(b, "add_nuw", a, c))
+        return NULL;
+    operands[0] = a;
+    operands[1] = c;
+    at = rocke_i_attrs(b);
+    rocke_attr_set_str(b, &at, "wrap", "nuw");
+    return rocke_i_op1(b, ROCKE_OP_ARITH_ADD, operands, 2, a->type, &at, "add");
+}
+
 rocke_value_t* rocke_b_sub(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c)
 {
     return rocke_i_binop(b, ROCKE_OP_ARITH_SUB, a, c, "sub");
@@ -436,6 +478,15 @@ rocke_value_t* rocke_b_cmp_ne(rocke_ir_builder_t* b, rocke_value_t* a, rocke_val
     return cmp_impl(b, a, c, "ne", "ne");
 }
 
+rocke_value_t* rocke_b_cmp_ult(rocke_ir_builder_t* b, rocke_value_t* a, rocke_value_t* c)
+{
+    if(!rocke_i_live(b))
+        return NULL;
+    if(!check_unsigned_int_operands(b, "cmp_ult", a, c))
+        return NULL;
+    return cmp_impl(b, a, c, "ult", "ult");
+}
+
 rocke_value_t*
     rocke_b_fcmp(rocke_ir_builder_t* b, const char* pred, rocke_value_t* a, rocke_value_t* c)
 {