EnzymeAD / EnzymeAD/Enzyme-JAX

DUSDUSPadPadToDUSPad

Open
#1,919 6 comments 1 reaction 2 assignees Claimed by @wsmoses View on GitHub
Dominant language
MLIR
Stars
131
Forks
53
Avg merge
1d 10h
Merged PRs (30d)
193

Description

```
%cst_161 = stablehlo.constant dense<0.000000e+00> : tensor loc(#loc)

%c_169 = stablehlo.constant dense<7> : tensor loc(#loc)
%c_171 = stablehlo.constant dense<8> : tensor loc(#loc)
%c_172 = stablehlo.constant dense<0> : tensor loc(#loc)

%503 = stablehlo.slice %iterArg_177 [8:12, 9:6136, 0:12272] : (tensor<20x6144x12272xf64>) -> tensor<4x6127x12272xf64> loc(#loc2288)

%504 = stablehlo.pad %503, %cst_161, low = [0, 1, 0], high = [0, 1, 0], interior = [0, 0, 0] : (tensor<4x6127x12272xf64>, tensor) -> tensor<4x6129x12272xf64> loc(#loc2288)
%505 = stablehlo.dynamic_update_slice %iterArg_177, %504, %c_171, %c_171, %c_172 : (tensor<20x6144x12272xf64>, tensor<4x6129x12272xf64>, tensor, tensor, tensor) -> tensor<20x6144x12272xf64> loc(#loc2288)

%506 = stablehlo.slice %505 [8:12, 7:6135, 0:12272] : (tensor<20x6144x12272xf64>) -> tensor<4x6128x12272xf64> loc(#loc2896)
%509 = stablehlo.slice %505 [8:12, 10:6138, 0:12272] : (tensor<20x6144x12272xf64>) -> tensor<4x6128x12272xf64> loc(#loc2896)

%2533 = stablehlo.slice %505 [8:12, 6:6135, 0:12272] : (tensor<20x6144x12272xf64>) -> tensor<4x6129x12272xf64> loc(#loc1390)
%2535 = stablehlo.slice %505 [8:12, 7:6136, 0:12272] : (tensor<20x6144x12272xf64>) -> tensor<4x6129x12272xf64> loc(#loc1390)
%2539 = stablehlo.slice %505 [8:12, 9:6138, 0:12272] : (tensor<20x6144x12272xf64>) -> tensor<4x6129x12272xf64> loc(#loc1390)

%512 = "enzymexla.extend"(%503) <{dimension = 0 : i64, lhs = 1 : i64, rhs = 1 : i64}> : (tensor<4x6127x12272xf64>) -> tensor<6x6127x12272xf64> loc(#loc3088)

%513 = stablehlo.pad %512, %cst_161, low = [0, 1, 0], high = [0, 0, 0], interior = [0, 0, 0] : (tensor<6x6127x12272xf64>, tensor) -> tensor<6x6128x12272xf64> loc(#loc3088)
%514 = stablehlo.dynamic_update_slice %505, %513, %c_169, %c_171, %c_172 : (tensor<20x6144x12272xf64>, tensor<6x6128x12272xf64>, tensor, tensor, tensor) -> tensor<20x6144x12272xf64> loc(#loc3088)

```

should be written similarly to https://github.com/EnzymeAD/Enzyme-JAX/pull/1911 (in this case here we can eliminate and/or fuse the two pads together, if possible.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.