diff --git a/coreai_torch/_aten_to_core.py b/coreai_torch/_aten_to_core.py index 74e18a6..75e596c 100644 --- a/coreai_torch/_aten_to_core.py +++ b/coreai_torch/_aten_to_core.py @@ -644,17 +644,20 @@ def replace_arange_start_step( else coreai.constant(1, dtype=start.type.element_type) ) - # Squeeze rank-1 inputs to 0D scalars for coreai.range_. + # coreai.range_ requires scalar (rank-0) operands that share an element + # type. aten.arange promotes mixed-type scalars internally; we replicate + # that here by squeezing each operand to rank-0 and casting to the FX + # node's output dtype before the op. + target_type = get_output_element_type_from_node(node) + def to_scalar(v: Value) -> Value: if v.type.rank > 0: - return coreai.shrink_dims(v, list(range(v.type.rank))) + v = coreai.shrink_dims(v, list(range(v.type.rank))) + if v.type.element_type != target_type: + v = coreai.cast(v, target_type) return v - result = coreai.range_(to_scalar(start), to_scalar(end), to_scalar(step)) - target_type = get_output_element_type_from_node(node) - if result.type.element_type != target_type: - result = coreai.cast(result, target_type) - return result + return coreai.range_(to_scalar(start), to_scalar(end), to_scalar(step)) def replace_batch_norm( @@ -814,6 +817,7 @@ def replace_binary_ops( "pow.Scalar": coreai.broadcasting_pow, "pow.Tensor_Tensor": coreai.broadcasting_pow, "pow.Tensor_Scalar": coreai.broadcasting_pow, + "pow": coreai.broadcasting_pow, "sub.Tensor": coreai.broadcasting_sub, "sub.Scalar": coreai.broadcasting_sub, "sub": coreai.broadcasting_sub, @@ -944,6 +948,36 @@ def replace_cat(values_map: dict[str, Value], node: fx.Node, loc: Location) -> V rank = inputs[0].type.rank dim = dim + rank if dim < 0 else dim + + # coreai.concat requires all non-concat dims to be provably equal across + # inputs. Under dynamic shapes, one branch can carry a dynamic non-concat + # axis while a sibling has a static size for the same axis — the dynamic + # side must in fact equal that static size, but the type system doesn't + # know it. Reshape such inputs to the known static size before the concat. + # Multiple distinct static sizes on one axis is a real mismatch and is + # left for the dialect verifier to reject. + dyn = ShapedType.get_dynamic_size() + + def known_static(axis: int) -> int | None: + if axis == dim: + return None + sizes = {inp.type.shape[axis] for inp in inputs if inp.type.shape[axis] != dyn} + return next(iter(sizes)) if len(sizes) == 1 else None + + statics = [known_static(a) for a in range(rank)] + promoted: list[Value] = [] + for inp in inputs: + new_shape = [ + statics[a] + if statics[a] is not None and inp.type.shape[a] == dyn + else inp.type.shape[a] + for a in range(rank) + ] + if new_shape != list(inp.type.shape): + inp = coreai.reshape(inp, new_shape) + promoted.append(inp) + inputs = promoted + return coreai.concat(dim, inputs) @@ -2595,6 +2629,7 @@ def replace_unary_ops( "log.default": coreai.log, "relu.default": coreai.relu, "round.default": coreai.round_, + "round": coreai.round_, "rsqrt.default": coreai.rsqrt, "sigmoid.default": coreai.sigmoid, "silu.default": coreai.silu, @@ -3434,6 +3469,7 @@ def sdpa_maskless(q: Value, k: Value, v: Value) -> Value: "pow.Scalar": replace_binary_ops, "pow.Tensor_Scalar": replace_binary_ops, "pow.Tensor_Tensor": replace_binary_ops, + "pow": replace_binary_ops, "prod.default": replace_prod_default, "prod.dim_int": replace_prod_dim_int, "reciprocal.default": replace_reciprocal, @@ -3441,6 +3477,7 @@ def sdpa_maskless(q: Value, k: Value, v: Value) -> Value: "remainder.Tensor": replace_remainder, "round.default": replace_unary_ops, "round.decimals": replace_round_decimals, + "round": replace_unary_ops, "repeat.default": replace_repeat, "rsqrt.default": replace_unary_ops, "scaled_dot_product_attention.default": replace_sdpa, diff --git a/coreai_torch/_utils.py b/coreai_torch/_utils.py index 813e36e..436034a 100644 --- a/coreai_torch/_utils.py +++ b/coreai_torch/_utils.py @@ -175,6 +175,23 @@ def __exit__(self, *exc: Any) -> None: self.close() +def to_rank1_int32(v: Value) -> Value: + """Coerce a SymInt-derived Value to canonical rank-1 si32 form. + + Dim-vector concats (used to build shape operands for ``coreai.reshape``, + ``coreai.interpolate``, etc.) require all inputs to share rank and + element type. SymInt values can arrive rank-0 (e.g. from + ``aten._local_scalar_dense``) or with a different int variant. This + helper produces the form that aligns with ``coreai.constant([i], + dtype=np.int32)`` and ``replace_sym_size_int``. + """ + if v.type.rank == 0: + v = coreai.reshape(v, [1]) + if v.type.element_type != IntegerType.get_signed(32): + v = coreai.cast(v, np.int32) + return v + + def upsample_build_output_shape_dynamic( x: Value, out_h: int | Value, out_w: int | Value ) -> Value: @@ -192,8 +209,8 @@ def upsample_build_output_shape_dynamic( ) shape = coreai.cast(coreai.get_shape(x), dtype=np.int32) non_spatial = coreai.slice_(shape, [0], [2], [1]) - h = [out_h] if isinstance(out_h, int) else out_h - w = [out_w] if isinstance(out_w, int) else out_w + h = [out_h] if isinstance(out_h, int) else to_rank1_int32(out_h) + w = [out_w] if isinstance(out_w, int) else to_rank1_int32(out_w) return coreai.concat(0, [non_spatial, h, w]) @@ -986,9 +1003,13 @@ def get_operand( if isinstance(arg, fx.Node): return values_map[arg.name] if isinstance(arg, list) and any(isinstance(e, fx.Node) for e in arg): - # Mixed list: resolve fx.Node elements via values_map, keep ints as constants. + # Mixed list: SymInt fx.Nodes + plain ints. Concat the two sources + # into a single rank-1 si32 dim vector. Both branches must produce + # the same canonical form so the concat verifier accepts them. dim_vals = [ - values_map[e.name] if isinstance(e, fx.Node) else coreai.constant([e]) + to_rank1_int32(values_map[e.name]) + if isinstance(e, fx.Node) + else coreai.constant([e], dtype=np.int32) for e in arg ] return coreai.concat(0, dim_vals) if len(dim_vals) > 1 else dim_vals[0] diff --git a/tests/ops/test_ops.py b/tests/ops/test_ops.py index ddfd4ac..a012e4a 100644 --- a/tests/ops/test_ops.py +++ b/tests/ops/test_ops.py @@ -347,6 +347,22 @@ def forward(self, x: Tensor) -> Tensor: x = torch.zeros(2, 8) await validate_numerical_output(model=model, x=x) + async def test_symint_end_with_float_start_step(self) -> None: + """Regression for ``replace_arange_start_step``: when ``end`` is + SymInt-derived (carrying f32 element type from a sym_size cast) + and ``start`` / ``step`` come in as scalar si32, ``coreai.range_`` + rejects the mismatched element types. The lowering must unify all + three to the FX-meta target type before the range op.""" + + class ArangeMixedModel(nn.Module): + def forward(self, x: Tensor) -> Tensor: + return torch.arange(0.0, x.shape[0], 0.5, dtype=torch.float32) + + model = ArangeMixedModel().eval() + x = torch.zeros(6) + dynamic_shapes = {"x": {0: torch.export.Dim("n", min=2, max=16)}} + await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes) + @pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize( @@ -985,6 +1001,29 @@ def forward(self, x: Tensor, empty: Tensor, y: Tensor) -> Tensor: """ filecheck_pattern(str(coreai_program), check_file=check_file) + async def test_dynamic_vs_static_non_concat_axis(self) -> None: + """Regression for non-concat-axis promotion: when one cat input has a + dynamic non-concat axis while a sibling has a known static size for + that same axis, the verifier can't statically prove they match. The + lowering reshapes the dynamic side to the known static size before + the concat (needed by multi-resolution feature merges under dynamic + shapes).""" + + class CatModel(nn.Module): + def forward(self, a: Tensor, b: Tensor) -> Tensor: + return torch.cat([a, b], dim=0) + + a = torch.rand(2, 4) + b = torch.rand(3, 4) + # Mark only a's dim 1 dynamic; b's dim 1 stays static at 4. Dim.AUTO + # prevents torch.export from specializing it back to a constant. + await validate_numerical_output( + model=CatModel().eval(), + a=a, + b=b, + dynamic_shapes=({1: torch.export.Dim.AUTO}, {}), + ) + @pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize( @@ -3221,6 +3260,142 @@ def forward(self, x: Tensor) -> Tensor: await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes) +class TestRound: + """Regression suite for the bare ``aten.round`` OpOverloadPacket target. + Same shape of bug as the pow case: torch.export rewrites can leave + ``aten.round`` without a ``.default`` overload, which previously raised + ``Unsupported ATen op: round`` / ``KeyError: 'round'``.""" + + @staticmethod + def _rewrite_to_overload_packet(program: object) -> None: + for node in program.graph_module.graph.nodes: + if node.op != "call_function": + continue + if getattr(node.target, "_overloadpacket", None) is torch.ops.aten.round: + node.target = torch.ops.aten.round + program.graph_module.recompile() + + def test_lowers_ir(self) -> None: + class RoundModel(nn.Module): + def forward(self, x: Tensor) -> Tensor: + return torch.round(x) + + x = torch.tensor([[-1.5, 2.0], [3.0, -4.0]]) + program = torch.export.export(RoundModel(), args=(x,)).run_decompositions() + self._rewrite_to_overload_packet(program) + + coreai_program = TorchConverter().add_exported_program(program).to_coreai() + coreai_program.optimize() + # The bare ``aten.round`` target must reach ``coreai.round`` — + # not produce a different op or fail with ``Unsupported ATen op``. + # Shape and element type must pass through unchanged. + filecheck_pattern( + str(coreai_program), + check_file=""" + // CHECK-LABEL: coreai.graph @main + // CHECK-SAME: %[[ARG0:.*]]: tensor<2x2xf32> + // CHECK-SAME: -> (tensor<2x2xf32> + // CHECK: %[[OUT:.+]] = coreai.round %[[ARG0]] : tensor<2x2xf32> -> tensor<2x2xf32> + // CHECK: coreai.output %[[OUT]] : tensor<2x2xf32> + """, + ) + + async def test_numerical(self) -> None: + class RoundModel(nn.Module): + def forward(self, x: Tensor) -> Tensor: + return torch.round(x) + + x = torch.tensor([[-1.5, 2.0, 0.4], [3.6, -4.5, -0.5]]) + program = torch.export.export(RoundModel(), args=(x,)).run_decompositions() + self._rewrite_to_overload_packet(program) + + coreai_program = TorchConverter().add_exported_program(program).to_coreai() + torch_out = RoundModel().eval()(x) + await validate_numerical_output( + coreai_program=coreai_program, torch_out=torch_out, x=x + ) + + +class TestView: + """Regression suite for ``get_operand``'s mixed-list path: a list arg + that mixes ``fx.Node`` SymInt entries with plain ints (e.g. a view + shape like ``[1, 1024, h, w]`` where h, w come from ``round(SymFloat)``) + must produce concat operands with a uniform rank-1 int32 element type + — otherwise the dialect rejects the dim-vector concat with ``expected + the same element type for all inputs to concat``.""" + + def test_view_with_round_symfloat_dims(self) -> None: + """View shape mixes plain ints with SymInts that come from + ``round(sym_float)`` shape arithmetic. Without the fix, + ``coreai.concat`` rejects the mixed-element-type operands when + building the dim vector for the reshape.""" + + class ViewModel(nn.Module): + def forward(self, x: Tensor) -> Tensor: + # round(SymFloat) → SymInt: produces a Value whose type + # doesn't match the int32 constants used for plain-int + # entries in the same dim list. + aspect = x.shape[3] / x.shape[2] + h = round((1800 / aspect) ** 0.5) + w = round((1800 * aspect) ** 0.5) + # ``view([1, C, round_h, round_w])`` — the mixed-list path + # in ``get_operand`` builds a dim-vector concat that + # combines plain-int constants with SymInt Values whose + # element type the fix normalises. + flat = x.view(1, 1024, -1) + return flat.view(1, 1024, h, w) + + x = torch.rand(1, 1024, 42, 42, dtype=torch.float16) + dynamic_shapes = { + "x": { + 2: torch.export.Dim("h", min=14), + 3: torch.export.Dim("w", min=14), + } + } + program = torch.export.export( + ViewModel().eval(), args=(x,), dynamic_shapes=dynamic_shapes + ).run_decompositions() + # Convert to Core AI — must not raise. Pre-fix this raised + # ``ValueError: Operation creation failed`` from ``coreai.concat`` + # with the diagnostic ``expected the same element type for all + # inputs to concat``. + coreai_program = TorchConverter().add_exported_program(program).to_coreai() + coreai_program.optimize() + # The fix's job is to normalise every entry of the view shape + # ``[1, 1024, h, w]`` to the same canonical rank-1 si32 form + # before the dim-vector concat. The check below pins: + # 1. The two SymInt entries (h, w) carry rank-1 f32 from the + # round(SymFloat) chain and get explicitly cast to rank-1 + # si32 (matching the int constants). + # 2. All four operands of the dim-vector concat are rank-1 + # si32, so the concat verifier accepts them. + # 3. The resulting rank-1 si32 vector of length 4 feeds the + # reshape that builds the final 4-D output. + filecheck_pattern( + str(coreai_program), + check_file=""" + // CHECK-LABEL: coreai.graph @main + // CHECK-SAME: %[[ARG0:.*]]: tensor<1x1024x?x?xf16> + // CHECK-SAME: -> (tensor<1x1024x?x?xf16> + // + // h and w land as rank-1 f32 from round(SymFloat) arithmetic; + // the fix casts each to rank-1 si32 so it matches the int + // constants in the same dim vector: + // CHECK: %[[H:.+]] = coreai.cast {{.*}} : tensor<1xf32> to tensor<1xsi32> + // CHECK: %[[W:.+]] = coreai.cast {{.*}} : tensor<1xf32> to tensor<1xsi32> + // + // Dim-vector concat: 4 entries [1, 1024, h, w] → all rank-1 si32: + // CHECK: %[[SHAPE:.+]] = coreai.concat {{.*}} : (tensor, tensor<1xsi32>, tensor<1xsi32>, tensor<1xsi32>, tensor<1xsi32>) -> tensor<4xsi32> + // + // Final reshape uses that shape vector to produce the 4-D output. + // (After optimize() the inner ``view(1, 1024, -1)`` is fused + // away, so the reshape applies directly to %arg0.): + // CHECK: %[[OUT:.+]] = coreai.reshape %[[ARG0]], %[[SHAPE]] : (tensor<1x1024x?x?xf16>, tensor<4xsi32>) -> tensor<1x1024x?x?xf16> + // CHECK: coreai.output %[[OUT]] + """, + ) + + _SCALAR_PARAMS = pytest.mark.parametrize( "scalar,tensor_shape,dtype", [ @@ -5546,6 +5721,66 @@ def forward(self, x: Tensor, target: Tensor) -> Tensor: model=model, x=x, target=target, dynamic_shapes=dynamic_shapes ) + def test_round_symfloat_size(self) -> None: + """Regression for ``upsample_build_output_shape_dynamic``: when + the output_size operands are SymInts that come from + ``round(SymFloat)`` shape arithmetic — e.g. + ``round((num_tokens / aspect_ratio) ** 0.5) * 14`` — the lowering + must reshape/cast each (out_h, out_w) operand to rank-1 int32 + before the concat that builds the output shape. Without the fix, + ``coreai.concat`` raises ``ValueError: Operation creation failed`` + on the mixed-rank / mixed-element-type inputs.""" + + class RoundSymFloatUpsample(nn.Module): + def forward(self, x: Tensor) -> Tensor: + aspect = x.shape[3] / x.shape[2] + h = round((1800 / aspect) ** 0.5) + w = round((1800 * aspect) ** 0.5) + return torch.nn.functional.interpolate( + x, size=(h * 14, w * 14), mode="nearest" + ) + + x = torch.rand(1, 3, 28, 28, dtype=torch.float16) + dynamic_shapes = { + "x": { + 2: torch.export.Dim("h", min=14), + 3: torch.export.Dim("w", min=14), + } + } + program = torch.export.export( + RoundSymFloatUpsample().eval(), args=(x,), dynamic_shapes=dynamic_shapes + ).run_decompositions() + # Convert must not raise. Pre-fix this raised ``ValueError: + # Operation creation failed`` from ``coreai.concat``. + coreai_program = TorchConverter().add_exported_program(program).to_coreai() + coreai_program.optimize() + # The output_shape concat must take three rank-1 si32 operands — + # the (N, C) slice from x's get_shape, plus normalised out_h and + # out_w. Pre-fix, out_h/out_w arrived as rank-1 f32 (from + # round(SymFloat) arithmetic) which broke the concat verifier. + filecheck_pattern( + str(coreai_program), + check_file=""" + // CHECK-LABEL: coreai.graph @main + // CHECK-SAME: %[[ARG0:.*]]: tensor<1x3x?x?xf16> + // + // out_h, out_w come in as rank-1 f32 from round-of-symfloat + // arithmetic; the fix casts each to rank-1 si32 to match the + // (N, C) slice of get_shape (which is also rank-1 si32): + // CHECK: %[[H:.+]] = coreai.cast {{.*}} : tensor<1xf32> to tensor<1xsi32> + // CHECK: %[[W:.+]] = coreai.cast {{.*}} : tensor<1xf32> to tensor<1xsi32> + // + // Output-shape concat: non_spatial (rank-1 si32, length 2) + // + out_h (rank-1 si32) + out_w (rank-1 si32) → rank-1 si32 length 4: + // CHECK: %[[OUT_SHAPE:.+]] = coreai.concat {{.*}} : (tensor, tensor<2xsi32>, tensor<1xsi32>, tensor<1xsi32>) -> tensor<4xsi32> + // + // The shape feeds nearest-neighbor interpolate (whole op + // matched on one line because the mode attribute is part of + // the same SSA statement): + // CHECK: coreai.interpolate {{.+}}, %[[OUT_SHAPE]], {{.+}} {interpolation_mode = #coreai.interpolation_mode} + """, + ) + class TestUpsampleBilinear2d: """Test suite for aten.upsample_bilinear2d.vec → coreai.interpolate (linear).""" @@ -5656,6 +5891,65 @@ def forward(self, x: Tensor, target: Tensor) -> Tensor: model=model, x=x, target=target, dynamic_shapes=dynamic_shapes ) + def test_round_symfloat_size(self) -> None: + """Regression for ``upsample_build_output_shape_dynamic``: when + the output_size operands are SymInts that come from + ``round(SymFloat)`` shape arithmetic — e.g. + ``round((num_tokens / aspect_ratio) ** 0.5) * 14`` — the lowering + must reshape/cast each (out_h, out_w) operand to rank-1 int32 + before the concat that builds the output shape. Without the fix, + ``coreai.concat`` raises ``ValueError: Operation creation failed`` + on the mixed-rank / mixed-element-type inputs.""" + + class RoundSymFloatUpsample(nn.Module): + def forward(self, x: Tensor) -> Tensor: + aspect = x.shape[3] / x.shape[2] + h = round((1800 / aspect) ** 0.5) + w = round((1800 * aspect) ** 0.5) + return torch.nn.functional.interpolate( + x, + size=(h * 14, w * 14), + mode="bilinear", + align_corners=False, + ) + + x = torch.rand(1, 3, 28, 28, dtype=torch.float16) + dynamic_shapes = { + "x": { + 2: torch.export.Dim("h", min=14), + 3: torch.export.Dim("w", min=14), + } + } + program = torch.export.export( + RoundSymFloatUpsample().eval(), args=(x,), dynamic_shapes=dynamic_shapes + ).run_decompositions() + # Convert must not raise. Pre-fix this raised ``ValueError: + # Operation creation failed`` from ``coreai.concat``. + coreai_program = TorchConverter().add_exported_program(program).to_coreai() + coreai_program.optimize() + # Same shape concat as the nearest case; only the interpolation + # mode tag differs. The fix is in the shape-build path so it + # applies identically here. + filecheck_pattern( + str(coreai_program), + check_file=""" + // CHECK-LABEL: coreai.graph @main + // CHECK-SAME: %[[ARG0:.*]]: tensor<1x3x?x?xf16> + // + // out_h, out_w cast from rank-1 f32 to rank-1 si32: + // CHECK: %[[H:.+]] = coreai.cast {{.*}} : tensor<1xf32> to tensor<1xsi32> + // CHECK: %[[W:.+]] = coreai.cast {{.*}} : tensor<1xf32> to tensor<1xsi32> + // + // Output-shape concat: all rank-1 si32 → rank-1 si32 length 4: + // CHECK: %[[OUT_SHAPE:.+]] = coreai.concat {{.*}} : (tensor, tensor<2xsi32>, tensor<1xsi32>, tensor<1xsi32>) -> tensor<4xsi32> + // + // The shape feeds bilinear interpolate (whole op matched on + // one line because the mode attribute is part of the same + // SSA statement): + // CHECK: coreai.interpolate {{.+}}, %[[OUT_SHAPE]], {{.+}} {interpolation_mode = #coreai.interpolation_mode} + """, + ) + @pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize( diff --git a/tests/ops/test_ops_ir.py b/tests/ops/test_ops_ir.py index 9ea8307..50c0786 100644 --- a/tests/ops/test_ops_ir.py +++ b/tests/ops/test_ops_ir.py @@ -994,25 +994,34 @@ def forward(self, x: Tensor) -> Tensor: x=x, dynamic_shapes={"x": {0: torch.export.Dim("batch", min=1)}}, ) - filecheck_pattern( - ir, - check_file=""" - // CHECK-LABEL: module { - // CHECK-NEXT: coreai.graph @main(%[[ARG0:.*]]: tensor {coreai.name = "x"}) -> (tensor {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} { - // CHECK-NEXT: %[[V0:.*]] = coreai.constant dense<> : tensor<0xui32> - // CHECK-NEXT: %[[V1:.*]] = coreai.constant dense<1> : tensor - // CHECK-NEXT: %[[V2:.*]] = coreai.constant dense<0> : tensor - // CHECK-NEXT: %[[V3:.*]] = coreai.constant dense<1> : tensor<1xsi32> - // CHECK-NEXT: %[[V4:.*]] = coreai.constant dense<0> : tensor<1xsi32> - // CHECK-NEXT: %[[V5:.*]] = coreai.get_shape %[[ARG0]] : tensor -> tensor<2xui32> - // CHECK-NEXT: %[[V6:.*]] = coreai.slice %[[V5]], %[[V4]], %[[V3]], %[[V3]] : (tensor<2xui32>, tensor<1xsi32>, tensor<1xsi32>, tensor<1xsi32>) -> tensor<1xui32> - // CHECK-NEXT: %[[V7:.*]] = coreai.cast %[[V6]] : tensor<1xui32> to tensor<1xsi32> - // CHECK-NEXT: %[[V8:.*]] = coreai.reshape %[[V7]], %[[V0]] : (tensor<1xsi32>, tensor<0xui32>) -> tensor - // CHECK-NEXT: %[[V9:.*]] = coreai.range %[[V2]], %[[V8]], %[[V1]] : (tensor, tensor, tensor) -> tensor - // CHECK-NEXT: %[[V10:.*]] = coreai.cast %[[V9]] : tensor to tensor - // CHECK-NEXT: coreai.output %[[V10]] : tensor - // CHECK-NEXT: } - // CHECK-NEXT: } + # ``end`` is a SymInt (rank-1 si32) sliced from x.shape[0]; + # ``start``/``step`` are scalar (rank-0) si32 constants. The + # lowering casts each operand to the FX node's output dtype (f32) + # before ``coreai.range_`` so the op sees uniform-typed scalars; + # the optimizer then constant-folds the casts on start/step into + # f32 constants directly. + filecheck_pattern( + ir, + check_file=""" + // CHECK-LABEL: coreai.graph @main + // CHECK-SAME: %[[ARG0:.*]]: tensor + // CHECK-SAME: -> (tensor + // + // start and step land as f32 constants directly (the si32 + // constant + cast-to-f32 pair gets folded by the optimizer): + // CHECK-DAG: %[[STEP:.+]] = coreai.constant dense<1.000000e+00> : tensor + // CHECK-DAG: %[[START:.+]] = coreai.constant dense<0.000000e+00> : tensor + // + // end: get_shape -> slice -> cast(ui32->si32) -> reshape(rank-1 to rank-0) -> cast(si32->f32): + // CHECK: %[[END_RANK1:.+]] = coreai.cast {{.*}} : tensor<1xui32> to tensor<1xsi32> + // CHECK: %[[END_RANK0:.+]] = coreai.reshape %[[END_RANK1]], {{.*}} : (tensor<1xsi32>, tensor<0xui32>) -> tensor + // CHECK: %[[END_F32:.+]] = coreai.cast %[[END_RANK0]] : tensor to tensor + // + // range called with all-f32 scalars; result is f32 directly, + // no post-range cast on the result: + // CHECK: %[[OUT:.+]] = coreai.range %[[START]], %[[END_F32]], %[[STEP]] : (tensor, tensor, tensor) -> tensor + // CHECK-NOT: coreai.cast %[[OUT]] + // CHECK: coreai.output %[[OUT]] : tensor """, )