Value binding methods when parsing assign statement. e.g. binding vi, vj, vk with T.axis.remap("SSR", [i, j, k]), when parsing vi, vj, vk = T.axis.remap("SSR", [i, j, k]). Parameters ---------- self : Parser The current parser. node : doc.expr The doc AS
(self: Parser, node: doc.expr, var_name: str, value: Any)
| 168 | |
| 169 | |
| 170 | def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) -> Any: |
| 171 | """Value binding methods when parsing assign statement. |
| 172 | e.g. binding vi, vj, vk with T.axis.remap("SSR", [i, j, k]), when parsing |
| 173 | vi, vj, vk = T.axis.remap("SSR", [i, j, k]). |
| 174 | |
| 175 | Parameters |
| 176 | ---------- |
| 177 | self : Parser |
| 178 | The current parser. |
| 179 | |
| 180 | node : doc.expr |
| 181 | The doc AST expression node for error reporting. |
| 182 | |
| 183 | var_name : str |
| 184 | The variable name. |
| 185 | |
| 186 | value : Any |
| 187 | The value to be bound with. |
| 188 | |
| 189 | Returns |
| 190 | ------- |
| 191 | res : Any |
| 192 | The bound value. |
| 193 | """ |
| 194 | if isinstance(value, T.scalar_wrapper): # pylint: disable=protected-access |
| 195 | # special case for scalar, name the buffer, but the var is used as BufferLoad |
| 196 | assert isinstance(value.scalar, T.BufferLoad) |
| 197 | IRBuilder.name(var_name, value.scalar.buffer) |
| 198 | return value.scalar |
| 199 | if isinstance(value, T.meta_var): |
| 200 | return value.value |
| 201 | elif getattr(type(value), "_is_meta_class", False): |
| 202 | name_meta_class_value(var_name, value) |
| 203 | return value |
| 204 | elif isinstance(value, list | tuple): |
| 205 | # Tuple-unpacking with a starred target (e.g. ``vi, *vs = T.axis.remap(...)``) |
| 206 | # collects multiple elements into a single list bound here. Recurse so each |
| 207 | # element gets a per-index name; this matches apache's behavior. |
| 208 | for i, v in enumerate(value): |
| 209 | bind_assign_value(self, node, f"{var_name}_{i}", v) |
| 210 | return value |
| 211 | elif isinstance(value, BufferRegion): |
| 212 | return value |
| 213 | elif isinstance(value, Frame): |
| 214 | value.add_callback(partial(value.__exit__, None, None, None)) |
| 215 | res = value.__enter__() |
| 216 | IRBuilder.name(var_name, res) |
| 217 | return res |
| 218 | elif isinstance(value, Buffer | IterVar | Layout) or ( |
| 219 | isinstance(value, Var) and not self.var_table.exist(value) |
| 220 | ): |
| 221 | IRBuilder.name(var_name, value) |
| 222 | return value |
| 223 | else: |
| 224 | if not isinstance(value, PrimExpr): |
| 225 | value = tvm.tirx.const(value) |
| 226 | if not isinstance(value, tvm.tirx.StringImm): |
| 227 | # x = expr -> scalar (auto-typed from value) |
nothing calls this directly
no test coverage detected
searching dependent graphs…