EnzymeAD / EnzymeAD/Enzyme-JAX

concatenate slice optimization

Open
#2,087 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
MLIR
Stars
131
Forks
53
Avg merge
1d 10h
Merged PRs (30d)
193

Description

```mlir
module @reactant_quantile attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<10x10xf64> {enzymexla.memory_effects = []}) -> tensor<5xf64> attributes {enzymexla.memory_effects = []} {
%c = stablehlo.constant dense<[[0], [24], [49], [74], [98]]> : tensor<5x1xi32>
%c_0 = stablehlo.constant dense<[[1], [25], [50], [75], [99]]> : tensor<5x1xi32>
%cst = stablehlo.constant dense<[0.000000e+00, 7.500000e-01, 5.000000e-01, 2.500000e-01, 1.000000e+00]> : tensor<5xf64>
%0 = stablehlo.negate %arg0 : tensor<10x10xf64>
%1 = stablehlo.reshape %0 : (tensor<10x10xf64>) -> tensor<100xf64>
%values, %indices = chlo.top_k(%1, k = 100) : tensor<100xf64> -> (tensor<100xf64>, tensor<100xi32>)
%2 = stablehlo.negate %values : tensor<100xf64>
%3 = stablehlo.reshape %2 : (tensor<100xf64>) -> tensor<100x1xf64>
%4 = stablehlo.slice %3 [0:1, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%5 = stablehlo.slice %3 [1:2, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%6 = stablehlo.slice %3 [2:3, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%7 = stablehlo.slice %3 [3:4, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%8 = stablehlo.slice %3 [4:5, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%9 = stablehlo.slice %3 [5:6, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%10 = stablehlo.slice %3 [6:7, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%11 = stablehlo.slice %3 [7:8, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%12 = stablehlo.slice %3 [8:9, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%13 = stablehlo.slice %3 [9:10, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%14 = stablehlo.concatenate %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x10xf64>
%15 = stablehlo.slice %3 [10:11, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%16 = stablehlo.slice %3 [11:12, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%17 = stablehlo.slice %3 [12:13, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%18 = stablehlo.slice %3 [13:14, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%19 = stablehlo.slice %3 [14:15, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%20 = stablehlo.slice %3 [15:16, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%21 = stablehlo.slice %3 [16:17, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%22 = stablehlo.slice %3 [17:18, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%23 = stablehlo.slice %3 [18:19, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%24 = stablehlo.slice %3 [19:20, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%25 = stablehlo.concatenate %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x10xf64>
%26 = stablehlo.slice %3 [20:21, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%27 = stablehlo.slice %3 [21:22, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%28 = stablehlo.slice %3 [22:23, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%29 = stablehlo.slice %3 [23:24, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%30 = stablehlo.slice %3 [24:25, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%31 = stablehlo.slice %3 [25:26, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%32 = stablehlo.slice %3 [26:27, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%33 = stablehlo.slice %3 [27:28, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%34 = stablehlo.slice %3 [28:29, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%35 = stablehlo.slice %3 [29:30, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%36 = stablehlo.concatenate %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x10xf64>
%37 = stablehlo.slice %3 [30:31, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%38 = stablehlo.slice %3 [31:32, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%39 = stablehlo.slice %3 [32:33, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%40 = stablehlo.slice %3 [33:34, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%41 = stablehlo.slice %3 [34:35, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%42 = stablehlo.slice %3 [35:36, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%43 = stablehlo.slice %3 [36:37, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%44 = stablehlo.slice %3 [37:38, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%45 = stablehlo.slice %3 [38:39, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%46 = stablehlo.slice %3 [39:40, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%47 = stablehlo.concatenate %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x10xf64>
%48 = stablehlo.slice %3 [40:41, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%49 = stablehlo.slice %3 [41:42, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%50 = stablehlo.slice %3 [42:43, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%51 = stablehlo.slice %3 [43:44, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%52 = stablehlo.slice %3 [44:45, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%53 = stablehlo.slice %3 [45:46, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%54 = stablehlo.slice %3 [46:47, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%55 = stablehlo.slice %3 [47:48, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%56 = stablehlo.slice %3 [48:49, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%57 = stablehlo.slice %3 [49:50, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%58 = stablehlo.concatenate %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x10xf64>
%59 = stablehlo.slice %3 [50:51, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%60 = stablehlo.slice %3 [51:52, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%61 = stablehlo.slice %3 [52:53, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%62 = stablehlo.slice %3 [53:54, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%63 = stablehlo.slice %3 [54:55, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%64 = stablehlo.slice %3 [55:56, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%65 = stablehlo.slice %3 [56:57, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%66 = stablehlo.slice %3 [57:58, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%67 = stablehlo.slice %3 [58:59, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%68 = stablehlo.slice %3 [59:60, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%69 = stablehlo.concatenate %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x10xf64>
%70 = stablehlo.slice %3 [60:61, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%71 = stablehlo.slice %3 [61:62, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%72 = stablehlo.slice %3 [62:63, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%73 = stablehlo.slice %3 [63:64, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%74 = stablehlo.slice %3 [64:65, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%75 = stablehlo.slice %3 [65:66, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%76 = stablehlo.slice %3 [66:67, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%77 = stablehlo.slice %3 [67:68, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%78 = stablehlo.slice %3 [68:69, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%79 = stablehlo.slice %3 [69:70, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%80 = stablehlo.concatenate %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x10xf64>
%81 = stablehlo.slice %3 [70:71, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%82 = stablehlo.slice %3 [71:72, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%83 = stablehlo.slice %3 [72:73, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%84 = stablehlo.slice %3 [73:74, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%85 = stablehlo.slice %3 [74:75, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%86 = stablehlo.slice %3 [75:76, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%87 = stablehlo.slice %3 [76:77, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%88 = stablehlo.slice %3 [77:78, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%89 = stablehlo.slice %3 [78:79, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%90 = stablehlo.slice %3 [79:80, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%91 = stablehlo.concatenate %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x10xf64>
%92 = stablehlo.slice %3 [80:81, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%93 = stablehlo.slice %3 [81:82, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%94 = stablehlo.slice %3 [82:83, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%95 = stablehlo.slice %3 [83:84, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%96 = stablehlo.slice %3 [84:85, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%97 = stablehlo.slice %3 [85:86, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%98 = stablehlo.slice %3 [86:87, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%99 = stablehlo.slice %3 [87:88, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%100 = stablehlo.slice %3 [88:89, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%101 = stablehlo.slice %3 [89:90, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%102 = stablehlo.concatenate %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x10xf64>
%103 = stablehlo.slice %3 [90:91, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%104 = stablehlo.slice %3 [91:92, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%105 = stablehlo.slice %3 [92:93, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%106 = stablehlo.slice %3 [93:94, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%107 = stablehlo.slice %3 [94:95, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%108 = stablehlo.slice %3 [95:96, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%109 = stablehlo.slice %3 [96:97, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%110 = stablehlo.slice %3 [97:98, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%111 = stablehlo.slice %3 [98:99, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%112 = stablehlo.slice %3 [99:100, 0:1] : (tensor<100x1xf64>) -> tensor<1x1xf64>
%113 = stablehlo.concatenate %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, dim = 1 : (tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>, tensor<1x1xf64>) -> tensor<1x10xf64>
%114 = stablehlo.concatenate %14, %25, %36, %47, %58, %69, %80, %91, %102, %113, dim = 0 : (tensor<1x10xf64>, tensor<1x10xf64>, tensor<1x10xf64>, tensor<1x10xf64>, tensor<1x10xf64>, tensor<1x10xf64>, tensor<1x10xf64>, tensor<1x10xf64>, tensor<1x10xf64>, tensor<1x10xf64>) -> tensor<10x10xf64>
%115 = stablehlo.reshape %114 : (tensor<10x10xf64>) -> tensor<100xf64>
%116 = "stablehlo.gather"(%115, %c) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<100xf64>, tensor<5x1xi32>) -> tensor<5xf64>
%117 = "stablehlo.gather"(%115, %c_0) <{dimension_numbers = #stablehlo.gather, indices_are_sorted = false, slice_sizes = array}> : (tensor<100xf64>, tensor<5x1xi32>) -> tensor<5xf64>
%118 = stablehlo.subtract %117, %116 : tensor<5xf64>
%119 = stablehlo.multiply %cst, %118 : tensor<5xf64>
%120 = stablehlo.add %116, %119 : tensor<5xf64>
return %120 : tensor<5xf64>
}
}
```

this can be written as slice -> transpose

Contributor guide

No contributing guide indexed for this repository

Research direction

The issue contains only a large MLIR module with repeated stablehlo.slice and stablehlo.concatenate operations; no repository file, test, or runnable entry point is named. First clarify the intended optimization and its correctness criteria, then locate the pass or regression-test area responsible for this IR and define a test showing the optimized result.

Written by the indexing model from the issue text.

Assessment

Domain
compilers, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
20/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.