Lines Matching full:shape
1 # 'shape' Dialect
3 Description of operations & types within the Shape dialect as well as their
4 [usage](#different-stages-of-lowering-shape-dialect).
8 ## Different stages of lowering Shape dialect
11 shape dialect and the lowering between these uses. Currently we have 3 worlds /
12 stages of lowering of shape functions:
15 This "input" form carries both the shape and whether in error state as
31 Starting from the shape function of matmul in the error monadic form
35 shape.function_library @shplib {
37 func.func @matmul(%lhs: !shape.value_shape, %rhs: !shape.value_shape) -> !shape.shape {
38 %c1 = shape.const_size 1
39 %c2 = shape.const_size 2
43 %lhs_shape = shape.shape_of %lhs : !shape.value_shape -> !shape.shape
44 %rhs_shape = shape.shape_of %rhs : !shape.value_shape -> !shape.shape
45 %lhs_rank = shape.rank %lhs_shape : !shape.shape -> !shape.size
46 %rhs_rank = shape.rank %rhs_shape : !shape.shape -> !shape.size
49 %r = "shape.meet"(%lhs_rank, %rhs_rank) : (!shape.size, !shape.size) -> !shape.size
50 %rank = shape.meet %c2, %r, error="requires rank 2 operands" :
51 !shape.size, !shape.size -> !shape.size
52 %l0, %l1 = "shape.split_at"(%lhs_shape, %c1) :
53 (!shape.shape, !shape.size) -> (!shape.shape, !shape.shape)
54 %r0, %r1 = "shape.split_at"(%rhs_shape, %c1) :
55 (!shape.shape, !shape.size) -> (!shape.shape, !shape.shape)
56 %c = shape.meet %l1, %r0, error="inner dimensions required to match" :
57 !shape.shape, !shape.shape -> !shape.shape
58 %res = shape.concat %l0, %r1
59 // Should have `shape.return %res requires %c, %rank` to enable
60 return %res : !shape.shape
69 ‘shape\_func’ as a special function op that allows passing multiple results
72 without handling the shape.return - that is a feature here as these are
75 avoid DCE until we have `shape.return`, at which point computing the
96 separate shape function library, while here we would normally reify it as part
97 of lowering, but for simplicity will show as a standalone shape function.
101 %c1 = shape.const_size 1
102 %c2 = shape.const_size 2
103 // We allow `shape.shape_of` to return either a `!shape.shape` or
106 %lhs_shape = shape.shape_of %lhs : tensor<*xf32> -> !shape.shape
107 %rhs_shape = shape.shape_of %rhs : tensor<*xf32> -> !shape.shape
108 %lhs_rank = shape.rank %lhs_shape : !shape.shape -> !shape.size
109 %rhs_rank = shape.rank %rhs_shape : !shape.shape -> !shape.size
110 %w1 = shape.cstr_eq %lhs_rank, %rhs_rank : !shape.witness
111 %res = shape.assuming %w1 -> tensor<?xindex> {
112 %r1 = shape.any %lhs_rank, %rhs_rank : (!shape.size, !shape.size) -> !shape.size
114 %w2 = shape.cstr_eq %c2, %r1, error="requires rank 2 operands"
115 %res_1 = shape.assuming %w2 -> tensor<?xindex> {
117 // %rank = shape.any %c2, %r1 (!shape.size, !shape.size) -> !shape.size
119 // then it could have been folded in `shape.any`.
120 %l0, %r0 = "shape.split_at"(%lhs_shape, %c1) :
121 (!shape.shape, !shape.size) -> !shape.shape
122 %l1, %r1 = "shape.split_at"(%lhs_shape, %c1) :
123 (!shape.shape, !shape.size) -> !shape.shape
124 %c = shape.meet %l1, %r0, error="inner dimensions required to match" :
125 !shape.size, !shape.size -> !shape.size
127 shape.assuming_yield %res
129 shape.assuming_yield %res_1
140 %c1 = shape.const_size 1
141 %c2 = shape.const_size 2
142 %lhs_shape = shape.shape_of %lhs : tensor<*xf32> -> tensor<?xindex>
143 %rhs_shape = shape.shape_of %rhs : tensor<*xf32> -> tensor<?xindex>
144 %lhs_rank = shape.rank %lhs_shape : tensor<?xindex> -> tensor<index>
145 %rhs_rank = shape.rank %rhs_shape : tensor<?xindex> -> tensor<index>
146 %w1 = shape.cstr_eq %c2, %lhs_rank, error="requires rank 2 operands"
147 %w2 = shape.cstr_eq %c2, %rhs_rank, error="requires rank 2 operands"
148 %w = shape.assuming_all %w1, %w2
149 %res = shape.assuming %w -> tensor<?xindex> {
150 %l0, %r0 = "shape.split_at"(%lhs_shape, %c1) :
151 (tensor<?xindex>, !shape.size) -> tensor<?xindex>
152 %l1, %r1 = "shape.split_at"(%lhs_shape, %c1) :
153 (tensor<?xindex>, !shape.size) -> tensor<?xindex>
154 %w3 = shape.cstr_eq %l1, %r0, error="inner dimensions required to match"
155 %res_2 = shape.assuming %w3 {
157 shape.assuming_yield %res
159 shape.assuming_yield %res_1
173 %lhs_shape = shape.shape_of %lhs : tensor<*xf32> -> tensor<?xindex>
174 %rhs_shape = shape.shape_of %rhs : tensor<*xf32> -> tensor<?xindex>
175 %lhs_rank = shape.rank %lhs_shape : tensor<?xindex> -> tensor<index>
176 %rhs_rank = shape.rank %rhs_shape : tensor<?xindex> -> tensor<index>
177 %w1 = shape.shape_eq %lhs_rank, %rhs_rank
178 %w2 = shape.shape_eq %c2, %lhs_rank
181 %l0, %l1 = shape.split_at(%lhs_shape, %c1) : tensor<?xindex>
182 %r0, %r1 = shape.split_at(%rhs_shape, %c1) : tensor<?xindex>
183 %w4 = shape.eq %l1, %r0
195 the `shape` dialect operations are no longer connected by producer-consumer