…degen (#176436)

Metal Shading Language rejects implicit float-to-bfloat conversions, so
bare float literals like `0.0` in generated shaders cause compilation
failures when the target variable is `bfloat` (or `half`). Three codegen
methods were affected:

- `constant()` ignored its `dtype` parameter and returned raw literals.
- `masked()` assigned a bare literal in the else-branch (`} else tmp = 0.0;`).
- `where()` passed a bare literal through the ternary without casting.

All three now emit `static_cast<bfloat>(...)` / `static_cast<half>(...)`
where needed. Tests added for half-precision constants, reductions, and
conditionals.

Pull Request resolved: #176436
Approved by: https://github.com/malfet

Test plan: Run `python -c "import torch;F=torch.nn.functional;print(torch.compile(lambda x: F.pad(F.gelu(x), [1, 0]))(torch.randn(4, device='mps', dtype=torch.bfloat16)))"`

(cherry picked from commit 3b161e7)

…degen (#176436) (#177193)

Metal Shading Language rejects implicit float-to-bfloat conversions, so
bare float literals like `0.0` in generated shaders cause compilation
failures when the target variable is `bfloat` (or `half`). Three codegen
methods were affected:

- `constant()` ignored its `dtype` parameter and returned raw literals.
- `masked()` assigned a bare literal in the else-branch (`} else tmp = 0.0;`).
- `where()` passed a bare literal through the ternary without casting.

All three now emit `static_cast<bfloat>(...)` / `static_cast<half>(...)`
where needed. Tests added for half-precision constants, reductions, and
conditionals.

Pull Request resolved: #176436
Approved by: https://github.com/malfet

Test plan: Run `python -c "import torch;F=torch.nn.functional;print(torch.compile(lambda x: F.pad(F.gelu(x), [1, 0]))(torch.randn(4, device='mps', dtype=torch.bfloat16)))"`

(cherry picked from commit 3b161e7)

Co-authored-by: Mergen Nachin <mnachin@meta.com>

- Remove `test_bfloat_constant`, `test_lowp_reduction`, and `test_lowp_where` as they don't test for anything beyond what existing tests cover.
- Add test_pad_after_gelu as a regression test for Voxtral compilation on MPS, exercising pad(gelu(x)) across fp32, fp16, and bfloat16.

Before #176436 test will fail with
```
torch._inductor.exc.InductorError: SyntaxError: failed to compile
    #include <c10/metal/utils.h>
    #include <c10/metal/special_math.h>
    kernel void generated_kernel(
        device bfloat* out_ptr0,
        constant bfloat* in_ptr0,
        uint xindex [[thread_position_in_grid]]
    ) {
        int x0 = (xindex) % (17);
        int x1 = c10::metal::floor_divide(xindex, 17);
        int x2 = xindex;
        auto tmp0 = (-1) + x0;
        auto tmp1 = static_cast<long>(tmp0);
        auto tmp2 = 0;
        auto tmp3 = tmp1 >= tmp2;
        bfloat tmp4;
        if (tmp3) {
            auto tmp_scoped_0 = static_cast<float>(in_ptr0[(-1) + x0 + 16*x1]);
            auto tmp_scoped_1 = static_cast<float>(tmp_scoped_0);
            auto tmp_scoped_2 = 0.5;
            auto tmp_scoped_3 = tmp_scoped_1 * tmp_scoped_2;
            auto tmp_scoped_4 = 0.7071067811865476;
            auto tmp_scoped_5 = tmp_scoped_1 * tmp_scoped_4;
            auto tmp_scoped_6 = c10::metal::erf(tmp_scoped_5);
            auto tmp_scoped_7 = 1.0;
            auto tmp_scoped_8 = tmp_scoped_6 + tmp_scoped_7;
            auto tmp_scoped_9 = tmp_scoped_3 * tmp_scoped_8;
            auto tmp_scoped_10 = static_cast<bfloat>(tmp_scoped_9);
            tmp4 = tmp_scoped_10;
        } else tmp4 = 0.0;
        out_ptr0[x2] = static_cast<bfloat>(tmp4);
    }
 with program_source:4495:23: error: assigning to 'bfloat' from incompatible type 'float'
        } else tmp4 = 0.0;
                      ^~~
```

Authored with Claude.


ghstack-source-id: 7919b53
Pull-Request: #177207

- Remove `test_bfloat_constant`, `test_lowp_reduction`, and `test_lowp_where` as they don't test for anything beyond what existing tests cover.
- Add test_pad_after_gelu as a regression test for Voxtral compilation on MPS, exercising pad(gelu(x)) across fp32, fp16, and bfloat16.

Before #176436 test will fail with
```
torch._inductor.exc.InductorError: SyntaxError: failed to compile
    #include <c10/metal/utils.h>
    #include <c10/metal/special_math.h>
    kernel void generated_kernel(
        device bfloat* out_ptr0,
        constant bfloat* in_ptr0,
        uint xindex [[thread_position_in_grid]]
    ) {
        int x0 = (xindex) % (17);
        int x1 = c10::metal::floor_divide(xindex, 17);
        int x2 = xindex;
        auto tmp0 = (-1) + x0;
        auto tmp1 = static_cast<long>(tmp0);
        auto tmp2 = 0;
        auto tmp3 = tmp1 >= tmp2;
        bfloat tmp4;
        if (tmp3) {
            auto tmp_scoped_0 = static_cast<float>(in_ptr0[(-1) + x0 + 16*x1]);
            auto tmp_scoped_1 = static_cast<float>(tmp_scoped_0);
            auto tmp_scoped_2 = 0.5;
            auto tmp_scoped_3 = tmp_scoped_1 * tmp_scoped_2;
            auto tmp_scoped_4 = 0.7071067811865476;
            auto tmp_scoped_5 = tmp_scoped_1 * tmp_scoped_4;
            auto tmp_scoped_6 = c10::metal::erf(tmp_scoped_5);
            auto tmp_scoped_7 = 1.0;
            auto tmp_scoped_8 = tmp_scoped_6 + tmp_scoped_7;
            auto tmp_scoped_9 = tmp_scoped_3 * tmp_scoped_8;
            auto tmp_scoped_10 = static_cast<bfloat>(tmp_scoped_9);
            tmp4 = tmp_scoped_10;
        } else tmp4 = 0.0;
        out_ptr0[x2] = static_cast<bfloat>(tmp4);
    }
 with program_source:4495:23: error: assigning to 'bfloat' from incompatible type 'float'
        } else tmp4 = 0.0;
                      ^~~
```

Authored with Claude.

ghstack-source-id: f075662
Pull-Request: #177207

…177207)

----

- Remove `test_bfloat_constant`, `test_lowp_reduction`, and `test_lowp_where` as they don't test for anything beyond what existing tests cover.
- Add `test_pad_after_gelu` as a regression test for Voxtral compilation on MPS, exercising pad(gelu(x)) across fp32, fp16, and bfloat16.

Before #176436 test will fail with
```
torch._inductor.exc.InductorError: SyntaxError: failed to compile
    #include <c10/metal/utils.h>
    #include <c10/metal/special_math.h>
    kernel void generated_kernel(
        device bfloat* out_ptr0,
        constant bfloat* in_ptr0,
        uint xindex [[thread_position_in_grid]]
    ) {
        int x0 = (xindex) % (17);
        int x1 = c10::metal::floor_divide(xindex, 17);
        int x2 = xindex;
        auto tmp0 = (-1) + x0;
        auto tmp1 = static_cast<long>(tmp0);
        auto tmp2 = 0;
        auto tmp3 = tmp1 >= tmp2;
        bfloat tmp4;
        if (tmp3) {
            auto tmp_scoped_0 = static_cast<float>(in_ptr0[(-1) + x0 + 16*x1]);
            auto tmp_scoped_1 = static_cast<float>(tmp_scoped_0);
            auto tmp_scoped_2 = 0.5;
            auto tmp_scoped_3 = tmp_scoped_1 * tmp_scoped_2;
            auto tmp_scoped_4 = 0.7071067811865476;
            auto tmp_scoped_5 = tmp_scoped_1 * tmp_scoped_4;
            auto tmp_scoped_6 = c10::metal::erf(tmp_scoped_5);
            auto tmp_scoped_7 = 1.0;
            auto tmp_scoped_8 = tmp_scoped_6 + tmp_scoped_7;
            auto tmp_scoped_9 = tmp_scoped_3 * tmp_scoped_8;
            auto tmp_scoped_10 = static_cast<bfloat>(tmp_scoped_9);
            tmp4 = tmp_scoped_10;
        } else tmp4 = 0.0;
        out_ptr0[x2] = static_cast<bfloat>(tmp4);
    }
 with program_source:4495:23: error: assigning to 'bfloat' from incompatible type 'float'
        } else tmp4 = 0.0;
                      ^~~
```

Authored with Claude.

Pull Request resolved: #177207
Approved by: https://github.com/atalman, https://github.com/mergennachin, https://github.com/jansel

…degen (pytorch#176436)

Metal Shading Language rejects implicit float-to-bfloat conversions, so
bare float literals like `0.0` in generated shaders cause compilation
failures when the target variable is `bfloat` (or `half`). Three codegen
methods were affected:

- `constant()` ignored its `dtype` parameter and returned raw literals.
- `masked()` assigned a bare literal in the else-branch (`} else tmp = 0.0;`).
- `where()` passed a bare literal through the ternary without casting.

All three now emit `static_cast<bfloat>(...)` / `static_cast<half>(...)`
where needed. Tests added for half-precision constants, reductions, and
conditionals.

Pull Request resolved: pytorch#176436
Approved by: https://github.com/malfet

…hader codegen (pytorch#176436)"

This reverts commit 4926192.

Reverted pytorch#176436 on behalf of https://github.com/zou3519 due to sorry I need to revert this in order to revert pytorch#176606 ([comment](pytorch#176436 (comment)))

…degen (pytorch#176436)

Metal Shading Language rejects implicit float-to-bfloat conversions, so
bare float literals like `0.0` in generated shaders cause compilation
failures when the target variable is `bfloat` (or `half`). Three codegen
methods were affected:

- `constant()` ignored its `dtype` parameter and returned raw literals.
- `masked()` assigned a bare literal in the else-branch (`} else tmp = 0.0;`).
- `where()` passed a bare literal through the ternary without casting.

All three now emit `static_cast<bfloat>(...)` / `static_cast<half>(...)`
where needed. Tests added for half-precision constants, reductions, and
conditionals.

Pull Request resolved: pytorch#176436
Approved by: https://github.com/malfet

…ytorch#177207)

----

- Remove `test_bfloat_constant`, `test_lowp_reduction`, and `test_lowp_where` as they don't test for anything beyond what existing tests cover.
- Add `test_pad_after_gelu` as a regression test for Voxtral compilation on MPS, exercising pad(gelu(x)) across fp32, fp16, and bfloat16.

Before pytorch#176436 test will fail with
```
torch._inductor.exc.InductorError: SyntaxError: failed to compile
    #include <c10/metal/utils.h>
    #include <c10/metal/special_math.h>
    kernel void generated_kernel(
        device bfloat* out_ptr0,
        constant bfloat* in_ptr0,
        uint xindex [[thread_position_in_grid]]
    ) {
        int x0 = (xindex) % (17);
        int x1 = c10::metal::floor_divide(xindex, 17);
        int x2 = xindex;
        auto tmp0 = (-1) + x0;
        auto tmp1 = static_cast<long>(tmp0);
        auto tmp2 = 0;
        auto tmp3 = tmp1 >= tmp2;
        bfloat tmp4;
        if (tmp3) {
            auto tmp_scoped_0 = static_cast<float>(in_ptr0[(-1) + x0 + 16*x1]);
            auto tmp_scoped_1 = static_cast<float>(tmp_scoped_0);
            auto tmp_scoped_2 = 0.5;
            auto tmp_scoped_3 = tmp_scoped_1 * tmp_scoped_2;
            auto tmp_scoped_4 = 0.7071067811865476;
            auto tmp_scoped_5 = tmp_scoped_1 * tmp_scoped_4;
            auto tmp_scoped_6 = c10::metal::erf(tmp_scoped_5);
            auto tmp_scoped_7 = 1.0;
            auto tmp_scoped_8 = tmp_scoped_6 + tmp_scoped_7;
            auto tmp_scoped_9 = tmp_scoped_3 * tmp_scoped_8;
            auto tmp_scoped_10 = static_cast<bfloat>(tmp_scoped_9);
            tmp4 = tmp_scoped_10;
        } else tmp4 = 0.0;
        out_ptr0[x2] = static_cast<bfloat>(tmp4);
    }
 with program_source:4495:23: error: assigning to 'bfloat' from incompatible type 'float'
        } else tmp4 = 0.0;
                      ^~~
```

Authored with Claude.

Pull Request resolved: pytorch#177207
Approved by: https://github.com/atalman, https://github.com/mergennachin, https://github.com/jansel