-
Notifications
You must be signed in to change notification settings - Fork 508
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[XLA:GPU][TRITON:XLA] Add ops and types to support TMA.
Tiled_tensor type and 3 ops: tile, insert, and extract to Triton_XLA dialect. These are going to be used to form common abstractions that would eventually lower to normal loads/stores or TMA variants. PiperOrigin-RevId: 726062943
- Loading branch information
1 parent
0c852dc
commit cc0fcd7
Showing
9 changed files
with
326 additions
and
21 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,43 @@ | ||
| /* Copyright 2025 The OpenXLA Authors. | ||
| Licensed under the Apache License, Version 2.0 (the "License"); | ||
| you may not use this file except in compliance with the License. | ||
| You may obtain a copy of the License at | ||
| http://www.apache.org/licenses/LICENSE-2.0 | ||
| Unless required by applicable law or agreed to in writing, software | ||
| distributed under the License is distributed on an "AS IS" BASIS, | ||
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| See the License for the specific language governing permissions and | ||
| limitations under the License. | ||
| ==============================================================================*/ | ||
|
|
||
| #include "llvm/ADT/TypeSwitch.h" // IWYU pragma: keep | ||
| #include "mlir/IR/DialectImplementation.h" // IWYU pragma: keep | ||
| #include "mlir/IR/OpImplementation.h" // IWYU pragma: keep | ||
| #include "xla/backends/gpu/codegen/triton/xla_triton_ops.h" | ||
|
|
||
| #define GET_ATTRDEF_CLASSES | ||
| #include "xla/backends/gpu/codegen/triton/xla_triton_attrs.cc.inc" | ||
| #define GET_TYPEDEF_CLASSES | ||
| #include "xla/backends/gpu/codegen/triton/xla_triton_types.cc.inc" | ||
|
|
||
| namespace mlir::triton::xla { | ||
|
|
||
| void XlaTritonDialect::initialize() { | ||
| addOperations< | ||
| #define GET_OP_LIST | ||
| #include "xla/backends/gpu/codegen/triton/xla_triton_ops.cc.inc" | ||
| >(); | ||
| addAttributes< | ||
| #define GET_ATTRDEF_LIST | ||
| #include "xla/backends/gpu/codegen/triton/xla_triton_attrs.cc.inc" | ||
| >(); | ||
| addTypes< | ||
| #define GET_TYPEDEF_LIST | ||
| #include "xla/backends/gpu/codegen/triton/xla_triton_types.cc.inc" | ||
| >(); | ||
| } | ||
|
|
||
| } // namespace mlir::triton::xla |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,42 @@ | ||
| /* Copyright 2025 The OpenXLA Authors. | ||
| Licensed under the Apache License, Version 2.0 (the "License"); | ||
| you may not use this file except in compliance with the License. | ||
| You may obtain a copy of the License at | ||
| http://www.apache.org/licenses/LICENSE-2.0 | ||
| Unless required by applicable law or agreed to in writing, software | ||
| distributed under the License is distributed on an "AS IS" BASIS, | ||
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| See the License for the specific language governing permissions and | ||
| limitations under the License. | ||
| ==============================================================================*/ | ||
|
|
||
| #include <cstdint> | ||
|
|
||
| #include "mlir/IR/OpImplementation.h" // IWYU pragma: keep | ||
| #include "mlir/IR/Types.h" // IWYU pragma: keep | ||
| #include "mlir/Support/LLVM.h" | ||
| #include "xla/backends/gpu/codegen/triton/xla_triton_ops.h" | ||
|
|
||
| namespace mlir::triton::xla { | ||
|
|
||
| mlir::Type TiledTensorType::parse(mlir::AsmParser &parser) { | ||
| mlir::SmallVector<int64_t, 4> shape; | ||
| mlir::Type type; | ||
| if (parser.parseLess() || | ||
| parser.parseDimensionList(shape, /*allowDynamic=*/false) || | ||
| parser.parseType(type) || parser.parseGreater()) { | ||
| return {}; | ||
| } | ||
| return TiledTensorType::get(parser.getContext(), shape, type); | ||
| } | ||
|
|
||
| void TiledTensorType::print(mlir::AsmPrinter &printer) const { | ||
| printer << "<"; | ||
| printer.printDimensionList(getShape()); | ||
| printer << "x" << getElementType() << ">"; | ||
| } | ||
|
|
||
| } // namespace mlir::triton::xla |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,59 @@ | ||
| /* Copyright 2025 The OpenXLA Authors. | ||
|
|
||
| Licensed under the Apache License, Version 2.0 (the "License"); | ||
| you may not use this file except in compliance with the License. | ||
| You may obtain a copy of the License at | ||
|
|
||
| http://www.apache.org/licenses/LICENSE-2.0 | ||
|
|
||
| Unless required by applicable law or agreed to in writing, software | ||
| distributed under the License is distributed on an "AS IS" BASIS, | ||
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| See the License for the specific language governing permissions and | ||
| limitations under the License. | ||
| ==============================================================================*/ | ||
|
|
||
| #ifndef XLA_BACKENDS_GPU_CODEGEN_TRITON_XLA_TRITON_TYPES_TD_ | ||
| #define XLA_BACKENDS_GPU_CODEGEN_TRITON_XLA_TRITON_TYPES_TD_ | ||
|
|
||
| include "xla/backends/gpu/codegen/triton/xla_triton_dialect.td" | ||
| include "mlir/IR/BuiltinTypes.td" // ValueSemantics | ||
| include "mlir/IR/BuiltinTypeInterfaces.td" | ||
|
|
||
| // ----------------------------------------------------------------------------- | ||
| // TiledTensorType | ||
| // ----------------------------------------------------------------------------- | ||
|
|
||
| class TTXLA_Type<string name, string typeMnemonic, list<Trait> traits = []> : | ||
| TypeDef<XlaTritonDialect, name, traits> { | ||
| let mnemonic = typeMnemonic; | ||
| } | ||
|
|
||
| def TTXLA_TiledTensorType : TTXLA_Type<"TiledTensor", "tiled_tensor", [ | ||
| ShapedTypeInterface, ValueSemantics]> { | ||
| let summary = "A tile of a tensor."; | ||
| let description = [{ | ||
| Usage: | ||
| This type will typically be constructed via triton_xla.tile op. The intent | ||
| is to capture tiling information and pass it along to other ops such as | ||
| triton_xla.extract and triton_xla.insert. Refer to the ops for examples. | ||
| }]; | ||
|
|
||
| let parameters = (ins | ||
| ArrayRefParameter<"int64_t">:$shape, | ||
| "Type":$elementType | ||
| ); | ||
|
|
||
| let hasCustomAssemblyFormat = 1; | ||
|
|
||
| let extraClassDeclaration = [{ | ||
| TiledTensorType cloneWith(std::optional<llvm::ArrayRef<int64_t>> shape, | ||
| mlir::Type elementType) const { | ||
| return TiledTensorType::get(getContext(), shape.value_or(getShape()), | ||
| elementType); | ||
| } | ||
| bool hasRank() const { return true; } | ||
| }]; | ||
| } | ||
|
|
||
| #endif // XLA_BACKENDS_GPU_CODEGEN_TRITON_XLA_TRITON_TYPES_TD_ |
Oops, something went wrong.